import numpy as np
import pandas as pd
from pathlib import Path
import itertools
import os
import sys
import torch
from torch.utils.data import Dataset, DataLoader
import transformers
import torch.nn as nn
import gc
from tqdm import tqdm

try:
    import text_extensions_for_pandas as tp
except ModuleNotFoundError as e:
        if ".." not in sys.path:
            sys.path.insert(0, "..")
        import text_extensions_for_pandas as tp


# init some globals we need later on 
BASE_DIR = 'Corpus/'
N_PREDICATES = 40
N_between_sample = 40
fold_paths = {"train": BASE_DIR + "CoNLL2009-ST-English-train.conllu",
              "eval": BASE_DIR + "CoNLL2009-ST-evaluation-English.conllu"}
conll_09_cols = ["LEMMA",'POS','pPOS','FEAT','head','DEPREL','FILLPRED','PRED']
bert_model_name = 'bert-base-uncased'
data_path = "Corpus"
MODEL_PATH = "Outputs"
train_ensemble = False



# class BertDumbedForTokenClassification(transformers.models.bert.modeling_bert.BertPreTrainedModel):

#     _keys_to_ignore_on_load_unexpected = [r"pooler"]

#     def __init__(self, config, n_comps_to_project=256,project_seed=42):
#         super().__init__(config)
#         self.num_labels = config.num_labels
#         self.n_comps_to_project = n_comps_to_project

#         self.bert = transformers.models.bert.modeling_bert.BertModel(config, add_pooling_layer=False)
#         self.dropout = nn.Dropout(config.hidden_dropout_prob)
#         self.proj_mat = self._make_proj_mat(config.hidden_size,self.n_comps_to_project,project_seed)
#         self.classifier = nn.Linear(self.n_comps_to_project, config.num_labels)
        

#     def forward(
#         self,
#         input_ids=None,
#         attention_mask=None,
#         token_type_ids=None,
#         position_ids=None,
#         head_mask=None,
#         inputs_embeds=None,
#         labels=None,
#         output_attentions=None,
#         output_hidden_states=None,
#         return_dict=None,
#     ):
#         r"""
#         labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
#             Labels for computing the token classification loss. Indices should be in ``[0, ..., config.num_labels -
#             1]``.
#         """
#         return_dict = return_dict if return_dict is not None else self.config.use_return_dict

#         outputs = self.bert(
#             input_ids,
#             attention_mask=attention_mask,
#             token_type_ids=token_type_ids,
#             position_ids=position_ids,
#             head_mask=head_mask,
#             inputs_embeds=inputs_embeds,
#             output_attentions=output_attentions,
#             output_hidden_states=output_hidden_states,
#             return_dict=return_dict,
#         )

#         sequence_output = outputs[0]
#         projected = self.proj_mat(sequence_output)
#         projected_dropped = self.dropout(projected)
#         logits = self.classifier(projected_dropped)

#         loss = None
#         if labels is not None:
#             loss_fct = nn.CrossEntropyLoss()
#             # Only keep active parts of the loss
#             if attention_mask is not None:
#                 active_loss = attention_mask.view(-1) == 1
#                 active_logits = logits.view(-1, self.num_labels)
#                 active_labels = torch.where(
#                     active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
#                 )
#                 loss = loss_fct(active_logits, active_labels)
#             else:
#                 loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))

#         if not return_dict:
#             output = (logits,) + outputs[2:]
#             return ((loss,) + output) if loss is not None else output

#         return transformers.modeling_outputs.TokenClassifierOutput(
#             loss=loss,
#             logits=logits,
#             hidden_states=outputs.hidden_states,
#             attentions=outputs.attentions,
#         )
    
#     def _make_proj_mat(self,hidden_size,projection_size,seed):
#         random_state = np.random.RandomState(seed)
#         projmat = torch.tensor(random_state.normal(loc=0,scale=1.0/np.sqrt(projection_size),size=(projection_size, hidden_size)))
#         layer = nn.Linear(hidden_size,projection_size)
#         layer.requires_grad=False
#         layer.weights = projmat
#         print(type(layer))
#         return layer




class BertDataset(Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        res = {}
        for key in self.encodings:
            res[key] = torch.tensor(self.encodings[key][idx])
        res['labels'] = torch.tensor(self.labels[idx])
        return res

    def __len__(self):
        return len(self.labels)
        

def make_predicates_dtype(values):
    dtype = pd.CategoricalDtype(values)
    int_to_cat = list(dtype.categories)
    cat_to_int = {cat:i for i, cat in enumerate(int_to_cat)}
    return dtype, int_to_cat, cat_to_int

# def re_align_labels(new_toks,labels):
#     NONE_VAL = 40
#     tokenized_labels = []
#     print_if = 0
#     for i,label in enumerate(labels):
#         mapping_list =new_toks.word_ids(i)
#         label_mappings = np.array(mapping_list[:sum((mapping is not None for mapping in mapping_list))])
#         doc_labels = np.ones(len(mapping_list),dtype=int)*NONE_VAL
#         try:
#             doc_labels[:len(label_mappings)] = np.array(label)[label_mappings]
#         except(IndexError): 
#             print(doc_labels)
#             print(repr(len(label_mappings)))
#             print(repr(np.array(label)))
#             print(repr(label_mappings))
#             print("label_mappings")
#             print("doc_labels")
#             print(repr(doc_labels[:len(label_mappings)]))
#             assert(False)
#         if print_if:
#             print("\nlabel")
#             print(label)
#             print(label_mappings)
#             print_if -=1
#         tokenized_labels.append(doc_labels.tolist())
#     return tokenized_labels

def encode_word_senses(encodings, ws):
    encoded_ws = []
    for i, (word_sense, offset) in enumerate(zip(ws, encodings.offset_mapping)):
        doc_enc_labels = np.ones(len(offset), dtype=int) * -100
        arr_offset = np.array(offset)
        doc_enc_labels[(arr_offset[:,0] == 0) & (arr_offset[:,1] != 0)] = word_sense

        encoded_ws.append(doc_enc_labels.tolist())
    return encoded_ws


# import documents from file. 
fold_docs = {}
for fold,fold_path in fold_paths.items(): 
    fold_docs[fold] = tp.io.conll.conll_u_to_dataframes(fold_path,column_names=conll_09_cols,separate_sentences_by_doc=True)
    print(f"converted fold: '{fold}' to list of {len(fold_docs[fold])} dataframes")

print(repr(fold_docs['train'][0].head(7)))

# create predicate names
train_preds_and_spans = pd.concat( (doc[['span','PRED']]     for doc in fold_docs['train'])  )
# drop predicate args because we won't be using them

# create a list of PRED arguments by most common 
common_fill_preds = (train_preds_and_spans
                     .rename(columns={'span':'instances'})
                     .groupby('PRED',dropna=True)
                     .count()
                     .sort_values('instances',ascending=False)
                    )

predicates_list = list(common_fill_preds.head(N_PREDICATES).index)
predicates_list.append('Misc')
predicates_list.append('X')


preds, int_to_pred, pred_to_int = make_predicates_dtype(predicates_list)


def conv_docs_to_dtype(fold,preds_list,pred_dtype,dtype_to_int):
    docs = []
    show = True
    for init_doc in fold:
        doc = init_doc.copy()
        doc.PRED.fillna('X',inplace=True)
        indexes = ~doc['PRED'].isin(preds_list)
        set_default = 'Misc'
        doc.loc[indexes,'PRED'] = set_default 
        doc['PRED_id'] = doc.PRED.apply(lambda x: dtype_to_int[x])
        doc.PRED = doc.PRED.astype(pred_dtype)
        
        docs.append(doc)
    return docs
        


train_fold_sents = conv_docs_to_dtype(fold_docs['train'],predicates_list,preds,pred_to_int)
eval_fold_sents = conv_docs_to_dtype(fold_docs['eval'],predicates_list,preds,pred_to_int)

# verify that we've done things correctly: 
print(repr(preds))
print(repr(int_to_pred))
print(len(int_to_pred))
print(repr(pred_to_int))


train_sents_words = [df.span.apply(lambda x: x.covered_text).to_list() 
                        for df in train_fold_sents]
train_sents_labels = [df.PRED_id.to_list() for df in train_fold_sents ]


# prepare for fine tuning of bert model
tokenizer = transformers.BertTokenizerFast.from_pretrained(bert_model_name, add_special_tokens=True)
#tokenize document 
tokenized_train = tokenizer(train_sents_words,return_offsets_mapping=True,
                            add_special_tokens=False,
                          padding=True, is_split_into_words=True, truncation=True)

tokenized_train_labels = encode_word_senses(tokenized_train,train_sents_labels)
train_data = BertDataset(tokenized_train, tokenized_train_labels)

print("show tokenized labels:")
for labs in tokenized_train_labels[:4]:
    print(labs,len(labs))


#tokenize eval portion
eval_sents_words = [df.span.apply(lambda x: x.covered_text).to_list() 
                        for df in eval_fold_sents]
eval_sents_labels = [df.PRED_id.to_list() for df in eval_fold_sents ]

tokenized_eval = tokenizer(eval_sents_words,return_offsets_mapping=True,
                            add_special_tokens=False,
                          padding=True, is_split_into_words=True, truncation=True)
tokenized_eval_labels = encode_word_senses(tokenized_eval,eval_sents_labels)
eval_data = BertDataset(tokenized_eval, tokenized_eval_labels)


## Create our specific dumbed model
if not train_ensemble:
    bert4tokclass = transformers.BertForTokenClassification.from_pretrained(bert_model_name,num_labels=len(int_to_pred))

    #get ready to train for real 
    train_loader = DataLoader(train_data, batch_size=64)
    eval_loader = DataLoader(eval_data, batch_size=32)
    optim = transformers.AdamW(bert4tokclass.parameters())
    device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

    print("training_model")
    bert4tokclass.to(device)
    bert4tokclass.train()
    all_train_losses = []
    all_eval_losses = []
    all_epochs = []
    results_on_eval = {
        'epoch': [],
        'precision': [],
        'recall': [],
        'f1': [],
    }
    i=0
    for epoch in tqdm(range(100)):
        train_loss = 0
        i=0
        for batch in train_loader:
            optim.zero_grad()
            outputs = bert4tokclass(
                input_ids=batch['input_ids'].to(device),
                attention_mask=batch['attention_mask'].to(device),
                token_type_ids=batch['token_type_ids'].to(device),
                labels=batch['labels'].to(device)
            )
            loss = outputs.loss
            loss.backward()
            optim.step()
            train_loss += loss.item()
            

            # get dev loss & perf
            if i% N_between_sample == 0:
                eval_loss = 0
                eval_preds = []
                with torch.no_grad():
                    bert4tokclass.eval()
                    for batch in eval_loader:
                        outputs = bert4tokclass(
                            input_ids=batch['input_ids'].to(device),
                            attention_mask=batch['attention_mask'].to(device),
                            token_type_ids=batch['token_type_ids'].to(device),
                            labels=batch['labels'].to(device)
                        )
                        eval_loss += outputs.loss.item()

                        # logits = [batch_size, seq_len, num_ws]
                        logits = outputs.logits
                        logits = torch.softmax(logits, dim=-1)

                        pred_unpadded = [list(itertools.compress(logits, labels != -100)) for logits, labels in zip(logits, batch['labels'])]
                        pred = [pred_logit for pred_logits in pred_unpadded for pred_logit in pred_logits]
                        eval_preds += pred        
                eval_loss = eval_loss / len(eval_loader)
                all_eval_losses.append(eval_loss)
                all_epochs.append(epoch)
                train_loss = train_loss / (train_loader.batch_size*N_between_sample)
                all_train_losses.append(train_loss)
                train_loss=0
            i +=1
        

        if not os.path.isdir(Path(data_path, f'model_full')):
            os.mkdir(Path(data_path, f'model_full'))
        losses = pd.DataFrame(data={'epoch': all_epochs, 'train_loss': all_train_losses,'eval_losses':all_eval_losses})
        losses.to_csv(Path(data_path,f'model_full', 'model_losses.csv'), index=False)


        if (epoch + 1) % 10 == 0:
            print(train_loss)
            torch.save(bert4tokclass.state_dict(), Path(data_path, f'model_full', 'epoch_{}_bert.pt'.format(epoch+1)))
        bert4tokclass.train()
        

    losses = pd.DataFrame(data={'epoch': list(range(1, len(all_train_losses)+1)), 'train_loss': all_train_losses})
    losses.to_csv(Path(data_path, 'finetuned_losses_final.csv'), index=False)
    exit()

# # if selected, train full ensemble

# compslist = [16,32,64,128,256]
# seedslist = [1,2,3]
# comp_seed_list = [(n_comps,seed) for n_comps in compslist for seed in seedslist]
# print(f'\n\nlist of components to seeds: {comp_seed_list}')
# print('\nTraining:')
# for comps,seed in comp_seed_list:
#     bert4tokclass = BertDumbedForTokenClassification.from_pretrained(bert_model_name,num_labels=len(int_to_pred),n_comps_to_project=comps,project_seed=seed)
#     bert = transformers.BertForTokenClassification.from_pretrained(bert_model_name, num_labels=len(int_to_pred))
#     bert.load_state_dict(torch.load(Path('Outputs', 'finetuned_bert', 'epoch_100_bert.pt'),map_location=torch.device('cpu')))
#     bert4tokclass.bert = bert.bert
#     bert4tokclass.bert.requires_grad = False
#     bert4tokclass.proj_mat.requires_grad = False
#     gc.collect()


#     #get ready to train for real 
#     train_loader = DataLoader(train_data, batch_size=32)
#     eval_loader = DataLoader(eval_data, batch_size=16)
#     optim = transformers.AdamW(bert4tokclass.classifier.parameters(), lr=1e-5)
#     device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

#     print("training_model")
#     best_loss = 1e9
#     bert4tokclass.to(device)
#     bert4tokclass.train()
#     all_train_losses = []
#     all_eval_losses = []
#     results_on_eval = {
#         'epoch': [],
#         'precision': [],
#         'recall': [],
#         'f1': [],
#     }

#     for epoch in tqdm(range(100)):
#         train_loss = 0
#         for batch in train_loader:
#             optim.zero_grad()
#             outputs = bert4tokclass(
#                 input_ids=batch['input_ids'].to(device),
#                 attention_mask=batch['attention_mask'].to(device),
#                 token_type_ids=batch['token_type_ids'].to(device),
#                 labels=batch['labels'].to(device)
#             )
#             loss = outputs.loss
#             loss.backward()
#             optim.step()
#             train_loss += loss.item()
#         train_loss = train_loss / len(train_loader)

#         # get dev loss & perf
#         eval_loss = 0
#         eval_preds = []
#         with torch.no_grad():
#             bert4tokclass.eval()
#             for batch in eval_loader:
#                 outputs = bert4tokclass(
#                     input_ids=batch['input_ids'].to(device),
#                     attention_mask=batch['attention_mask'].to(device),
#                     token_type_ids=batch['token_type_ids'].to(device),
#                     labels=batch['labels'].to(device)
#                 )
#                 eval_loss += outputs.loss.item()

#                 # logits = [batch_size, seq_len, num_ws]
#                 logits = outputs.logits
#                 logits = torch.softmax(logits, dim=-1)

#                 pred_unpadded = [list(itertools.compress(logits, labels != -100)) for logits, labels in zip(logits, batch['labels'])]
#                 pred = [pred_logit for pred_logits in pred_unpadded for pred_logit in pred_logits]
#                 eval_preds += pred        
#         eval_loss = eval_loss / len(eval_loader)

        
#         all_train_losses.append(train_loss)
#         all_eval_losses.append(eval_loss)
        

#         if not os.path.isdir(Path(data_path, f'model_dummbed_{comps}_{seed}')):
#             os.mkdir(Path(data_path, f'model_dummbed_{comps}_{seed}'))
#         losses = pd.DataFrame(data={'epoch': list(range(1, len(all_train_losses)+1)), 'train_loss': all_train_losses,'eval_losses':all_eval_losses})
#         losses.to_csv(Path(data_path,f'model_dummbed_{comps}_{seed}', 'model_losses.csv'), index=False)


#         if (epoch + 1) % 20 == 0:
#             print(train_loss)
#             torch.save(bert4tokclass.state_dict(), Path(data_path, f'model_dummbed_{comps}_{seed}', 'epoch_{}_bert.pt'.format(epoch+1)))
#         bert4tokclass.train()
        

#     losses = pd.DataFrame(data={'epoch': list(range(1, len(all_train_losses)+1)), 'train_loss': all_train_losses})
#     losses.to_csv(Path(data_path, 'finetuned_losses.csv'), index=False)
