import logging
import os
import time
from typing import Any

from huggingface_hub import PyTorchModelHubMixin
from pytorch_lightning import Trainer, LightningModule, LightningDataModule
from pytorch_lightning.utilities.types import OptimizerLRScheduler, STEP_OUTPUT, EVAL_DATALOADERS
from torch.utils.data import DataLoader, Dataset, IterableDataset
from torcheval.metrics import BinaryAccuracy, BinaryAUROC, BinaryF1Score, BinaryPrecision, BinaryRecall
from transformers import BertModel, BatchEncoding, BertTokenizer, TrainingArguments
from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions
import torch
from torch import nn
from datasets import load_dataset
timber = logging.getLogger()
# logging.basicConfig(level=logging.DEBUG)
logging.basicConfig(level=logging.INFO)  # change to level=logging.DEBUG to print more logs...

NO_REGULARIZATION = 0
L1_REGULARIZATION_CODE = 1
L2_REGULARIZATION_CODE = 2
L1_AND_L2_REGULARIZATION_CODE = 3

black = "\u001b[30m"
red = "\u001b[31m"
green = "\u001b[32m"
yellow = "\u001b[33m"
blue = "\u001b[34m"
magenta = "\u001b[35m"
cyan = "\u001b[36m"
white = "\u001b[37m"

FORWARD = "FORWARD_INPUT"
BACKWARD = "BACKWARD_INPUT"

DNA_BERT_6 = "zhihan1996/DNA_bert_6"


class CommonAttentionLayer(nn.Module):
  def __init__(self, hidden_size, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.attention_linear = nn.Linear(hidden_size, 1)
    pass

  def forward(self, hidden_states):
    # Apply linear layer
    attn_weights = self.attention_linear(hidden_states)
    # Apply softmax to get attention scores
    attn_weights = torch.softmax(attn_weights, dim=1)
    # Apply attention weights to hidden states
    context_vector = torch.sum(attn_weights * hidden_states, dim=1)
    return context_vector, attn_weights


class ReshapedBCEWithLogitsLoss(nn.BCEWithLogitsLoss):
  def forward(self, input, target):
    return super().forward(input.squeeze(), target.float())


class MQtlDnaBERT6Classifier(nn.Module, PyTorchModelHubMixin):
  def __init__(self,
               bert_model=BertModel.from_pretrained(pretrained_model_name_or_path=DNA_BERT_6),
               hidden_size=768,
               num_classes=1,
               *args,
               **kwargs
               ):
    super().__init__(*args, **kwargs)

    self.model_name = "MQtlDnaBERT6Classifier"

    self.bert_model = bert_model
    self.attention = CommonAttentionLayer(hidden_size)
    self.classifier = nn.Linear(hidden_size, num_classes)
    pass

  def forward(self, input_ids: torch.tensor, attention_mask: torch.tensor, token_type_ids):
    """
    # torch.Size([128, 1, 512]) --> [128, 512]
    input_ids = input_ids.squeeze(dim=1).to(DEVICE)
    # torch.Size([16, 1, 512]) --> [16, 512]
    attention_mask = attention_mask.squeeze(dim=1).to(DEVICE)
    token_type_ids = token_type_ids.squeeze(dim=1).to(DEVICE)
    """
    bert_output: BaseModelOutputWithPoolingAndCrossAttentions = self.bert_model(
      input_ids=input_ids,
      attention_mask=attention_mask,
      token_type_ids=token_type_ids
    )

    last_hidden_state = bert_output.last_hidden_state
    context_vector, ignore_attention_weight = self.attention(last_hidden_state)
    y = self.classifier(context_vector)
    return y

"""
class TorchMetrics:
  def __init__(self):
    self.binary_accuracy = BinaryAccuracy()  #.to(device)
    self.binary_auc = BinaryAUROC()  # .to(device)
    self.binary_f1_score = BinaryF1Score()  # .to(device)
    self.binary_precision = BinaryPrecision()  # .to(device)
    self.binary_recall = BinaryRecall()  # .to(device)
    pass

  def update_on_each_step(self, batch_predicted_labels, batch_actual_labels):  # todo: Add log if needed
    # it looks like the library maintainers changed preds to input, ie, before: preds, now: input
    self.binary_accuracy.update(input=batch_predicted_labels, target=batch_actual_labels)
    self.binary_auc.update(input=batch_predicted_labels, target=batch_actual_labels)
    self.binary_f1_score.update(input=batch_predicted_labels, target=batch_actual_labels)
    self.binary_precision.update(input=batch_predicted_labels, target=batch_actual_labels)
    self.binary_recall.update(input=batch_predicted_labels, target=batch_actual_labels)
    pass

  def compute_and_log_on_each_step(self, log, log_prefix: str, log_color: str = green):
    b_accuracy = self.binary_accuracy.compute()
    b_auc = self.binary_auc.compute()
    b_f1_score = self.binary_f1_score.compute()
    b_precision = self.binary_precision.compute()
    b_recall = self.binary_recall.compute()
    timber.info(log_color + f"{log_prefix}_acc = {b_accuracy}, {log_prefix}_auc = {b_auc}, {log_prefix}_f1_score = {b_f1_score}, {log_prefix}_precision = {b_precision}, {log_prefix}_recall = {b_recall}")
    log(f"{log_prefix}_accuracy", b_accuracy)
    log(f"{log_prefix}_auc", b_auc)
    log(f"{log_prefix}_f1_score", b_f1_score)
    log(f"{log_prefix}_precision", b_precision)
    log(f"{log_prefix}_recall", b_recall)

  # def reset_on_epoch_end(self):
  #   self.binary_accuracy.reset()
  #   self.binary_auc.reset()
  #   self.binary_f1_score.reset()
  #   self.binary_precision.reset()
  #   self.binary_recall.reset()

  def compute_and_reset_on_epoch_end(self, log, log_prefix: str, log_color: str = green):
    b_accuracy = self.binary_accuracy.compute()
    b_auc = self.binary_auc.compute()
    b_f1_score = self.binary_f1_score.compute()
    b_precision = self.binary_precision.compute()
    b_recall = self.binary_recall.compute()
    timber.info(  log_color + f"{log_prefix}_acc = {b_accuracy}, {log_prefix}_auc = {b_auc}, {log_prefix}_f1_score = {b_f1_score}, {log_prefix}_precision = {b_precision}, {log_prefix}_recall = {b_recall}")
    log(f"{log_prefix}_accuracy", b_accuracy)
    log(f"{log_prefix}_auc", b_auc)
    log(f"{log_prefix}_f1_score", b_f1_score)
    log(f"{log_prefix}_precision", b_precision)
    log(f"{log_prefix}_recall", b_recall)

    self.binary_accuracy.reset()
    self.binary_auc.reset()
    self.binary_f1_score.reset()
    self.binary_precision.reset()
    self.binary_recall.reset()
    pass
"""


class TorchMetrics:
  def __init__(self):
    self.binary_accuracy = BinaryAccuracy()  #.to(device)
    self.binary_auc = BinaryAUROC()  # .to(device)
    self.binary_f1_score = BinaryF1Score()  # .to(device)
    self.binary_precision = BinaryPrecision()  # .to(device)
    self.binary_recall = BinaryRecall()  # .to(device)
    pass

  def update_on_each_step(self, batch_predicted_labels, batch_actual_labels):  # todo: Add log if needed
    # it looks like the library maintainers changed preds to input, ie, before: preds, now: input
    self.binary_accuracy.update(input=batch_predicted_labels, target=batch_actual_labels)
    self.binary_auc.update(input=batch_predicted_labels, target=batch_actual_labels)
    self.binary_f1_score.update(input=batch_predicted_labels, target=batch_actual_labels)
    self.binary_precision.update(input=batch_predicted_labels, target=batch_actual_labels)
    self.binary_recall.update(input=batch_predicted_labels, target=batch_actual_labels)
    pass

  def compute_metrics_and_log(self, log, log_prefix: str, log_color: str = green):
    b_accuracy = self.binary_accuracy.compute()
    b_auc = self.binary_auc.compute()
    b_f1_score = self.binary_f1_score.compute()
    b_precision = self.binary_precision.compute()
    b_recall = self.binary_recall.compute()
    timber.info(  log_color + f"{log_prefix}_acc = {b_accuracy}, {log_prefix}_auc = {b_auc}, {log_prefix}_f1_score = {b_f1_score}, {log_prefix}_precision = {b_precision}, {log_prefix}_recall = {b_recall}")
    log(f"{log_prefix}_accuracy", b_accuracy)
    log(f"{log_prefix}_auc", b_auc)
    log(f"{log_prefix}_f1_score", b_f1_score)
    log(f"{log_prefix}_precision", b_precision)
    log(f"{log_prefix}_recall", b_recall)

    pass

  def reset_on_epoch_end(self):
    self.binary_accuracy.reset()
    self.binary_auc.reset()
    self.binary_f1_score.reset()
    self.binary_precision.reset()
    self.binary_recall.reset()


class MQtlBertClassifierLightningModule(LightningModule):
  def __init__(self,
               classifier: nn.Module,
               criterion=None,  # nn.BCEWithLogitsLoss(),
               regularization: int = L2_REGULARIZATION_CODE,  # 1 == L1, 2 == L2, 3 (== 1 | 2) == both l1 and l2, else ignore / don't care
               l1_lambda=0.0001,
               l2_wright_decay=0.0001,
               *args: Any,
               **kwargs: Any):
    super().__init__(*args, **kwargs)
    self.classifier = classifier
    self.criterion = criterion
    self.train_metrics = TorchMetrics()
    self.validate_metrics = TorchMetrics()
    self.test_metrics = TorchMetrics()

    self.regularization = regularization
    self.l1_lambda = l1_lambda
    self.l2_weight_decay = l2_wright_decay
    pass

  def forward(self, x, *args: Any, **kwargs: Any) -> Any:
    input_ids: torch.tensor = x["input_ids"]
    attention_mask: torch.tensor = x["attention_mask"]
    token_type_ids: torch.tensor = x["token_type_ids"]
    # print(f"\n{ type(input_ids) = }, {input_ids = }")
    # print(f"{ type(attention_mask) = }, { attention_mask = }")
    # print(f"{ type(token_type_ids) = }, { token_type_ids = }")

    return self.classifier.forward(input_ids, attention_mask, token_type_ids)

  def configure_optimizers(self) -> OptimizerLRScheduler:
    # Here we add weight decay (L2 regularization) to the optimizer
    weight_decay = 0.0
    if self.regularization == 2 or self.regularization == 3:
      weight_decay = self.l2_weight_decay
    return torch.optim.Adam(self.parameters(), lr=1e-5, weight_decay=weight_decay)  # , weight_decay=0.005)

  def training_step(self, batch, batch_idx, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
    # Accuracy on training batch data
    x, y = batch
    preds = self.forward(x)
    loss = self.criterion(preds, y)

    if self.regularization == 1 or self.regularization == 3:  # apply l1 regularization
      l1_norm = sum(p.abs().sum() for p in self.parameters())
      loss += self.l1_lambda * l1_norm

    self.log("train_loss", loss)
    # calculate the scores start
    self.train_metrics.update_on_each_step(batch_predicted_labels=preds.squeeze(), batch_actual_labels=y)
    self.train_metrics.compute_metrics_and_log(log=self.log, log_prefix="train")
    # self.train_metrics.compute_and_log_on_each_step(log=self.log, log_prefix="train")
    # calculate the scores end
    return loss

  def on_train_epoch_end(self) -> None:
    self.train_metrics.compute_metrics_and_log(log=self.log, log_prefix="train")
    self.train_metrics.reset_on_epoch_end()
    pass

  def validation_step(self, batch, batch_idx, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
    # Accuracy on validation batch data
    # print(f"debug { batch = }")
    x, y = batch
    preds = self.forward(x)
    loss = self.criterion(preds, y)
    """ loss = 0 #  <------------------------- maybe the loss calculation is problematic """
    self.log("valid_loss", loss)
    # calculate the scores start
    self.validate_metrics.update_on_each_step(batch_predicted_labels=preds.squeeze(), batch_actual_labels=y)
    self.validate_metrics.compute_metrics_and_log(log=self.log, log_prefix="validate", log_color=blue)
    # self.validate_metrics.compute_and_log_on_each_step(log=self.log, log_prefix="validate", log_color=blue)

    # calculate the scores end
    return loss

  def on_validation_epoch_end(self) -> None:
    self.validate_metrics.compute_metrics_and_log(log=self.log, log_prefix="validate", log_color=blue)
    self.validate_metrics.reset_on_epoch_end()
    return None

  def test_step(self, batch, batch_idx, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
    # Accuracy on validation batch data
    x, y = batch
    preds = self.forward(x)
    loss = self.criterion(preds, y)
    self.log("test_loss", loss)  # do we need this?
    # calculate the scores start
    self.test_metrics.update_on_each_step(batch_predicted_labels=preds.squeeze(), batch_actual_labels=y)
    self.test_metrics.compute_metrics_and_log(log=self.log, log_prefix="test", log_color=magenta)
    # self.test_metrics.compute_and_log_on_each_step(log=self.log, log_prefix="test", log_color=magenta)

    # calculate the scores end
    return loss

  def on_test_epoch_end(self) -> None:
    self.test_metrics.compute_metrics_and_log(log=self.log, log_prefix="test", log_color=magenta)
    self.test_metrics.reset_on_epoch_end()
    return None

  pass


class PagingMQTLDnaBertDataset(IterableDataset):
  def __init__(self, dataset, tokenizer, max_length=512):  # hold on! why is it 512? I added 4000, and it crashed, the error suggested 512, that's why 512
    self.dataset = dataset
    self.bert_tokenizer = tokenizer
    self.max_length = max_length

  # def __len__(self):
  #   return len(self.dataset)

  def __iter__(self):
    for row in self.dataset:
      processed = self.preprocess(row)
      if processed is not None:
        yield processed

  def preprocess(self, row):
    sequence = row['sequence']  # Fetch the 'sequence' column
    label = row['label']  # Fetch the 'label' column (or whatever target you use)

    # Tokenize the sequence
    encoded_sequence: BatchEncoding = self.bert_tokenizer(
      sequence,
      truncation=True,
      padding='max_length',
      max_length=self.max_length,
      return_tensors='pt'
    )

    encoded_sequence_squeezed = {key: val.squeeze() for key, val in encoded_sequence.items()}
    return encoded_sequence_squeezed, label


class DNABERTDataModule(LightningDataModule):
  def __init__(self, model_name=DNA_BERT_6, batch_size=8, WINDOW=-1, is_local=False):
    super().__init__()
    self.tokenized_dataset = None
    self.dataset = None
    self.train_dataset: PagingMQTLDnaBertDataset = None
    self.validate_dataset: PagingMQTLDnaBertDataset = None
    self.test_dataset: PagingMQTLDnaBertDataset = None
    self.tokenizer = BertTokenizer.from_pretrained(pretrained_model_name_or_path=model_name)
    self.batch_size = batch_size
    self.is_local = is_local
    self.window = WINDOW

  def prepare_data(self):
    # Download and prepare dataset
    data_files = {
      # small samples
      "train_binned_200": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_200_train_binned.csv",
      "validate_binned_200": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_200_validate_binned.csv",
      "test_binned_200": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_200_test_binned.csv",
      # medium samples
      "train_binned_1000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_1000_train_binned.csv",
      "validate_binned_1000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_1000_validate_binned.csv",
      "test_binned_1000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_1000_test_binned.csv",

      # large samples
      "train_binned_4000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_train_binned.csv",
      "validate_binned_4000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_validate_binned.csv",
      "test_binned_4000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_test_binned.csv",

      # really tiny
      # "tiny_train": "/home/soumic/Codes/mqtl-classification/src/inputdata/tiny_dataset_4000_train_binned.csv",
      # "tiny_validate": "/home/soumic/Codes/mqtl-classification/src/inputdata/tiny_dataset_4000_validate_binned.csv",
      # "tiny_test": "/home/soumic/Codes/mqtl-classification/src/inputdata/tiny_dataset_4000_test_binned.csv",

      "tiny_train": "/home/soumic/Codes/mqtl-classification/src/inputdata/medium_dataset_4000_train_binned.csv",
      "tiny_validate": "/home/soumic/Codes/mqtl-classification/src/inputdata/medium_dataset_4000_validate_binned.csv",
      "tiny_test": "/home/soumic/Codes/mqtl-classification/src/inputdata/medium_dataset_4000_test_binned.csv",
    }
    if self.is_local:
      self.dataset = load_dataset("csv", data_files=data_files, streaming=True)
    else:
      self.dataset = load_dataset("fahimfarhan/mqtl-classification-datasets")

  def setup(self, stage=None):
    self.train_dataset = PagingMQTLDnaBertDataset(self.dataset['tiny_train'], self.tokenizer)
    self.validate_dataset = PagingMQTLDnaBertDataset(self.dataset['tiny_validate'], self.tokenizer)
    self.test_dataset = PagingMQTLDnaBertDataset(self.dataset['tiny_test'], self.tokenizer)

  def train_dataloader(self):
    return DataLoader(self.train_dataset, batch_size=self.batch_size, num_workers=1)

  def val_dataloader(self):
    return DataLoader(self.validate_dataset, batch_size=self.batch_size, num_workers=1)

  def test_dataloader(self) -> EVAL_DATALOADERS:
    return DataLoader(self.test_dataset, batch_size=self.batch_size, num_workers=1)


def start_bert(classifier_model, model_save_path, criterion, WINDOW, batch_size=4,
                is_binned=True, is_debug=False, max_epochs=10, regularization_code = L2_REGULARIZATION_CODE):
  is_my_laptop = os.path.isfile("/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_train_binned.csv")
  model_local_directory = f"my-awesome-model-{WINDOW}"
  model_remote_repository = f"fahimfarhan/dnabert-6-mqtl-classifier-{WINDOW}"
  file_suffix = ""
  if is_binned:
    file_suffix = "_binned"

  data_module = DNABERTDataModule(batch_size=batch_size, WINDOW=WINDOW, is_local=is_my_laptop)

  # classifier_model = classifier_model.to(DEVICE)

  classifier_module = MQtlBertClassifierLightningModule(
    classifier=classifier_model,
    regularization=regularization_code, criterion=criterion)

  # if os.path.exists(model_save_path):
  #   classifier_module.load_state_dict(torch.load(model_save_path))

  classifier_module = classifier_module  # .double()

  # Prepare data using the DataModule
  data_module.prepare_data()
  data_module.setup()

  trainer = Trainer(max_epochs=max_epochs, precision="32")

  # Train the model
  trainer.fit(model=classifier_module, datamodule=data_module)
  trainer.test(model=classifier_module, datamodule=data_module)
  torch.save(classifier_module.state_dict(), model_save_path)

  # classifier_module.push_to_hub("fahimfarhan/mqtl-classifier-model")

  classifier_model.save_pretrained(save_directory=model_local_directory, safe_serialization=False)
  # push to the hub
  commit_message = f":tada: Push model for window size {WINDOW} from huggingface space"
  if is_my_laptop:
    commit_message = f":tada: Push model for window size {WINDOW} from zephyrus"

  classifier_model.push_to_hub(
    repo_id=model_remote_repository,
    # subfolder=f"my-awesome-model-{WINDOW}", subfolder didn't work :/
    commit_message=commit_message,  # f":tada: Push model for window size {WINDOW}"
    # safe_serialization=False
  )
  pass


if __name__ == "__main__":
  start_time = time.time()

  dataset_folder_prefix = "inputdata/"
  pytorch_model = MQtlDnaBERT6Classifier()
  start_bert(classifier_model=pytorch_model, model_save_path=f"weights_{pytorch_model.model_name}.pth",
             criterion=ReshapedBCEWithLogitsLoss(), WINDOW=4000, batch_size=12,  # max 14 on my laptop...
            max_epochs=1, regularization_code=L2_REGULARIZATION_CODE)

  # Record the end time
  end_time = time.time()
  # Calculate the duration
  duration = end_time - start_time
  # Print the runtime
  print(f"Runtime: {duration:.2f} seconds")

  pass