import time
from bert import *
from fastformer import *
from sys import getsizeof

bert_leaves = [
    "BertSelfAttention-3","BertSelfOutput-3", \
    "BertIntermedate-2","BertOutput-2", \
    "PositionalEmbedding-1", "Poolers-1", \
    "Embedding-0", "Classifier-0"
]

fastformer_leaves = [
    "FastSelfAttention-3","BertSelfOutput-3", \
    "BertIntermedate-2","BertOutput-2", \
    "PositionalEmbedding-1", "Poolers-1", \
    "Embedding-0", "Classifier-0"
]

ff_self_attention_record_functions = [
    "ff_fc",
    "ff_alpha",
    "ff_global_query",
    "ff_qk_interaction",
    "ff_beta",
    "ff_global_key",
    "ff_kv_interaction",
    "FastSelfAttention-3"
]

bert_record_functions = bert_leaves + ["BertformerEncoder-0", "BertformerLayer-1", "BertAttention-2", "Bert"]
fastformer_record_functions = fastformer_leaves  + ["FastformerEncoder-0", "FastformerLayer-1", "FastAttention-2", "Fastformer"]


def get_prof(news_words, label, train_index, test_index, model_name="bert", batch=1, max_seq_len=128, with_cuda=True, dummy=False):
    model = get_model(word_dict=word_dict, model_name=model_name, max_seq_len=max_seq_len, with_cuda=with_cuda)
    prof = test_iter (model, news_words, label, train_index, test_index, batch=batch, max_seq_len=max_seq_len, with_cuda=with_cuda, dummy=dummy)
    
    return prof

def get_stats (evt, parent, tab="|", verbose=True, df=None, level=0):

    sum_cpu_time=0
    sum_cuda_time=0
    sum_flops=0

    is_leaf = evt.cpu_children == []

    prefix     = "{}└─{}".format(tab, evt.name)
    cpu_time   = "CPU time(self) {} ({}) us".format(int(evt.cpu_time_total), int(evt.self_cpu_time_total))
    cuda_time  = "CUDA time(self) {} ({}) us".format(int(evt.cuda_time_total), int(evt.self_cuda_time_total))
    cpu_mem    = "CPU memory {} ({})".format(int(evt.cpu_memory_usage), int(evt.self_cpu_memory_usage))
    cuda_mem   = "CUDA memory {} ({})".format(int(evt.cuda_memory_usage), int(evt.self_cuda_memory_usage))
    
    print_string = "{0:_<30}, {1:_<30}, {2:_<30}, {3:_<30}, {4:_<40}, {5:_<10} FLOPS, parent {6:_<30}".format( prefix, \
                                                                                        cpu_time, cuda_time, \
                                                                                        cpu_mem, cuda_mem, evt.flops, parent)
    if verbose :
        if evt.flops != 0:
            print(termcolor.colored(print_string, "red"))
        elif evt.cpu_memory_usage != 0:
            print(termcolor.colored(print_string, "green"))
        elif (evt.name == "FastSelfAttention-3") or (evt.name == "BertSelfAttention-3"):
            print(termcolor.colored(print_string, "blue"))
        else :
            print(print_string)
    
#     print(evt.cpu_children)

    # Append current event to df list
    to_append = [evt.name, 
                int(evt.cpu_time_total),     int(evt.self_cpu_time_total), \
                int(evt.cuda_time_total),    int(evt.self_cuda_time_total), \
                int(evt.cpu_memory_usage),   int(evt.self_cpu_memory_usage), \
                int(evt.cuda_memory_usage),  int(evt.self_cuda_memory_usage), \
                evt.flops, level, is_leaf, parent]
    df.append(to_append)


    # Add current event
    sum_cpu_time += evt.self_cpu_time_total
    sum_cuda_time+= evt.self_cuda_time_total
    sum_flops    += evt.flops

    # Add child event 
    for child_evt in evt.cpu_children:

        child_cpu_time, child_flops = get_stats(child_evt, parent=parent, tab=tab+"|", verbose=verbose, df=df, level=level+1)

        sum_cpu_time += child_cpu_time
        sum_flops    += child_flops

    return (sum_cpu_time, sum_flops)


def get_func_tree (prof, verbose=True, target=None):
    funct_evts = prof.function_events
    
    cols = ["name", 
            "cpu_time_total",     "self_cpu_time_total",
            "cuda_time_total",    "self_cuda_time_total",
            "cpu_memory_usage",   "self_cpu_memory_usage",
            "cuda_memory_usage",  "self_cuda_memory_usage", 
            "flops", "level", "is_leaf", "parent"]

    df_master = []

    for evt in funct_evts :

        if evt.name == "[memory]":
            continue

        # Top-level Nodes
        if target is None :
            if evt.cpu_parent is None:
                cpu_time, flops = get_stats(evt, parent=evt.name, verbose=verbose, tab="", df=df_master)
                if verbose :
                    print()
                    
        elif evt.name in target:
            cpu_time, flops = get_stats(evt, parent=evt.name, verbose=verbose, tab="", df=df_master)
            if verbose :
                print()
                
    df_master = pd.DataFrame(df_master, columns=cols)
    
    return df_master


def plot_memory(df_master, record_func, with_cuda=True, L2_size=3e6, L1_size= 48e3):
    
    df = df_master[~df_master["name"].isin(record_func)].reset_index().reset_index()
    df = df.astype({"level_0": str})
    df['x'] = df["name"] + "-" + df["level_0"]
    
    plt.rcParams['figure.figsize'] = [40, 6]
#     plt.yscale("log")
    
    fig = plt.Figure()
    fig = sns.barplot(data=df, x="x", y="cuda_memory_usage" if with_cuda else "cpu_memory_usage")
    fig.tick_params(labelrotation=90)
    plt.yticks(fontsize=15)
    plt.ylabel("CUDA Memory Usage [B]", fontsize=15)
    
    for func in record_func :
        fig = plt.axvline(x=func, color="blue",  linestyle="dashed")
    
    # Cache
#     fig = plt.axhline(y=L1_size, color="green", linestyle="dashed")
#     fig = plt.axhline(y=L2_size, color="blue",  linestyle="dashed")
    fig = plt.axhline(y=L1_size + L2_size, color="red",  linestyle="dashed")
        
    plt.show()
    
    
def plot_memory_cumulative(df_master, record_func, with_cuda=True, L2_size=3e6, L1_size= 48e3):
    
    df = df_master[~df_master["name"].isin(record_func)].reset_index()
    df.index.name = 'kernel_idx'
    df.reset_index(inplace=True)
    df = df.astype({"kernel_idx": str})
    
    df['kernel'] = df["name"] + "-" + df["kernel_idx"]
    
    plt.rcParams['figure.figsize'] = [40, 6]
#     plt.yscale("log")
    
    fig = plt.Figure()
    fig = sns.barplot(data=df, x="kernel", y="cumulative_cuda_mem" if with_cuda else "cpu_memory_usage")
    fig.tick_params(labelrotation=90)
    plt.yticks(fontsize=15)
    plt.ylabel("CUDA Memory Usage [B]", fontsize=15)
    
    # Cache
#     fig = plt.axhline(y=L1_size, color="green", linestyle="dashed")
#     fig = plt.axhline(y=L2_size, color="blue",  linestyle="dashed")
    fig = plt.axhline(y=L1_size + L2_size, color="red",  linestyle="dashed")
        
    plt.show()


def preprocess_data (text, label, max_seq_len=256, dummy=False) :
    
    word_dict={'PADDING':0}
    
    if dummy :
        news_words=None
        label=None
        train_index=None
        test_index=None
        
    else :
        for sent in text:    
            for token in sent:        
                if token not in word_dict:
                    word_dict[token]=len(word_dict)

        news_words = []
        
        for sent in text:       
            sample=[]
            for token in sent:     
                sample.append(word_dict[token])
            sample = sample[:max_seq_len]
            news_words.append(sample+[0]*(max_seq_len-len(sample)))
        
        news_words=np.array(news_words,dtype='int32') 
        label=np.array(label,dtype='int32')
    
        index=np.arange(len(label))
        train_index=index[:120000]
        np.random.shuffle(train_index)
        test_index=index[120000:]
    
    return news_words, label, train_index, test_index, word_dict


def load_data(data_name='ag_news'):
    dataset = load_dataset(data_name)
    text=[]
    label=[]
    for row in dataset['train']['text']+dataset['test']['text']:
        text.append(wordpunct_tokenize(row.lower()))
    for row in dataset['train']['label']+dataset['test']['label']:
        label.append(row)
    
    return text, label
    
def get_model(word_dict, model_name="bert", max_seq_len=256, with_cuda=True):
    config=BertConfig.from_json_file("base_config.json")
    
    config.max_position_embeddings = max_seq_len

    if model_name == "fastformer":
        model = FastformerModel(config, word_dict)
    elif model_name == "bert":
        model = BertModel(config, word_dict)
    else :
        print("ERROR : No such model name")
    
    return model

def acc(y_true, y_hat):
    y_hat = torch.argmax(y_hat, dim=-1)
    tot = y_true.shape[0]
    hit = torch.sum(y_true == y_hat)
    return hit.data.float() * 1.0 / tot

# inference
def test_iter (model, news_words, label, train_index, test_index, batch=64, max_seq_len=256, with_cuda=True, dummy=False):
    
    # Configure Data
    if dummy :
        log_ids=np.zeros((batch,max_seq_len))
        targets = np.zeros((batch))
    else :
        log_ids=news_words[train_index][cnt*batch:cnt*batch+batch,:max_seq_len]
        targets= label[train_index][cnt*batch:cnt*batch+batch]

    if with_cuda :
        model.cuda()
        log_ids_ = torch.LongTensor(log_ids).cuda(non_blocking=True)
        targets_ = torch.LongTensor(targets).cuda(non_blocking=True)
    else :
        log_ids_ = torch.LongTensor(log_ids)
        targets_ = torch.LongTensor(targets)

    # Warmup
    for cnt in range(16):
        bz_loss, y_hat = model(log_ids_, targets_)
        torch.cuda.synchronize()
        
#     print("Memory allocated : {} B\nMemory reserved : {} B\nMemory cached : {} B\nSummary:\n{}".format(
#         torch.cuda.memory_allocated(),
#         torch.cuda.memory_reserved(),
#         torch.cuda.memory_cached(),
#         torch.cuda.memory_summary())
#      )      
            
    print("log_ids ({}) : {}, targets ({}) : {}".format(log_ids.shape, getsizeof(log_ids), targets.shape, getsizeof(targets)))
        
    del log_ids_, targets_
            
    # Test Iteration
    with profiler.profile(with_stack=True, profile_memory=True, with_flops=True, record_shapes=True, use_cuda=with_cuda) as prof:

        with profiler.record_function("Data Allocation"):
            if with_cuda :
                log_ids_ = torch.LongTensor(log_ids).cuda(non_blocking=True)
                targets_ = torch.LongTensor(targets).cuda(non_blocking=True)
            else :
                log_ids_ = torch.LongTensor(log_ids)
                targets_ = torch.LongTensor(targets)

        with profiler.record_function("Model Allocation"):
            if with_cuda : 
                model.cuda()
        
        with profiler.record_function("Inference"):
            bz_loss, y_hat = model(log_ids_, targets_)
            torch.cuda.synchronize()

    time.sleep(2)

    
#     print("Memory allocated : {} B\nMemory reserved : {} B\nMemory cached : {} B\nSummary:\n{}".format(
#         torch.cuda.memory_allocated(),
#         torch.cuda.memory_reserved(),
#         torch.cuda.memory_cached(),
#         torch.cuda.memory_summary())
#      )        
    
    
    if with_cuda :
        del log_ids, targets
        
    torch.cuda.empty_cache()
        
    return prof

# train
def train_iter (model, news_words, label, train_index, test_index, batch=64, max_seq_len=256, with_cuda=True, dummy=False):
    optimizer = optim.Adam([ {'params': model.parameters(), 'lr': 1e-3}])
    for epoch in range(2):
        loss = 0.0
        accuary = 0.0
        
        # GPU warmup
        for i in range(10):
            
            if dummy :
                log_ids = np.zeros((batch,max_seq_len))
                targets = np.zeros((batch))
            else :
                log_ids=news_words[train_index][cnt*batch:cnt*batch+batch,:max_seq_len]
                targets= label[train_index][cnt*batch:cnt*batch+batch]

            if with_cuda :
                log_ids = torch.LongTensor(log_ids).cuda(non_blocking=True)
                targets = torch.LongTensor(targets).cuda(non_blocking=True)
            else :
                log_ids = torch.LongTensor(log_ids)
                targets = torch.LongTensor(targets)
                
            bz_loss, y_hat = model(log_ids, targets)
            
        # Training
        for cnt in range(10):

            if dummy :
                log_ids=np.zeros((batch,max_seq_len))
                targets = np.zeros((batch))
            else :
                log_ids=news_words[train_index][cnt*batch:cnt*batch+batch,:max_seq_len]
                targets= label[train_index][cnt*batch:cnt*batch+batch]

            if with_cuda :
                log_ids = torch.LongTensor(log_ids).cuda(non_blocking=True)
                targets = torch.LongTensor(targets).cuda(non_blocking=True)
            else :
                log_ids = torch.LongTensor(log_ids)
                targets = torch.LongTensor(targets)

            with profiler.profile(with_stack=True, profile_memory=True, with_flops=True, record_shapes=True, use_cuda=with_cuda) as prof:
                    
                bz_loss, y_hat = model(log_ids, targets)
                loss += bz_loss.data.float()
                accuary += acc(targets, y_hat)
                unified_loss=bz_loss
                optimizer.zero_grad()
                unified_loss.backward()
                optimizer.step()
                
                torch.cuda.synchronize()
                
            if with_cuda :
                del log_ids, targets
                
        print( ' Ed: {}, train_loss: {:.5f}, acc: {:.5f}'.format(cnt * 64, loss.data / (cnt+1), accuary / (cnt+1)))
        
        torch.cuda.empty_cache()
    
    return prof

