import logging
import os
import time

import torch
import wandb
from easydict import EasyDict
from sklearn.metrics import precision_recall_fscore_support
from terminaltables import AsciiTable
from torch.optim import AdamW
from torch.optim.lr_scheduler import LinearLR
from tqdm import tqdm

from sdsvkie.datasets.base_dataset import BaseDataset
from sdsvkie.models.layoutlmv2 import LayoutLMv2
from sdsvkie.models.layoutlm import LayoutLM
from sdsvkie.utils import yaml_save, get_info_env, get_logger
from transformers import get_scheduler

# logging = logging.getlogging(__name__)
# logging.basicConfig(level=logging.INFO)



class Trainer:
    def __init__(self, cfg: dict):
        self.cfg = EasyDict(cfg)
        self.model = None
        self.processor = None
        self._init_trainer()

    def _init_trainer(self):
        
        if self.cfg.v3:
            os.environ["TOKENIZERS_PARALLELISM"] = "false"

        if not os.path.exists(self.cfg.save_dir):
            os.makedirs(self.cfg.save_dir, exist_ok=True)

        timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
        log_file = "{}/{}.log".format(self.cfg.save_dir, str(timestamp))
        self.logger = get_logger(log_file=log_file, log_level='INFO')
        self.logger.info(self.cfg)
        # self.logger.info(get_info_env())

        if not torch.cuda.is_available():
            self.logger.info("Can not found cuda, training with CPU!!!")
            self.cfg.device = "cpu"

        
        
        self.model = LayoutLMv2._load_model(self.cfg)
        self.processor = LayoutLMv2._load_processor(self.cfg)

        self.model.to(self.cfg.device)

        if self.cfg.wandb:
            wandb.init(
                project=self.cfg.wandb,
            )

    def _build_dataloader(self, data_dir, cache_file, use_sampling=False):
        base_dataset = BaseDataset(self.cfg)
        dataloader = base_dataset.build_dataloader_from_dir(
            data_dir=data_dir,
            processor=self.processor,
            device="cpu",
            batch_size=self.cfg.batch_size,
            num_workers=self.cfg.num_workers,
            shuffle=False,
            cache_file=cache_file,
            use_sampling=use_sampling,
        )
        return dataloader

    def val(self, val_dir=None):
        val_dir = val_dir if val_dir is not None else self.cfg.val_dir
        val_cache_file = os.path.splitext(val_dir)[0] + ".pkl"
        val_dataloader = self._build_dataloader(
            val_dir if val_dir is not None else self.cfg.val_dir,
            cache_file=val_cache_file
        )
        acc = self.val_on_dataloader(val_dataloader)
        return acc

    def train(self, train_dir=None, val_dir=None):
        self.logger.info("Building train dataloader...")
        base_dataset = BaseDataset(self.cfg)
        train_dir = train_dir if train_dir is not None else self.cfg.train_dir
        val_dir = val_dir if val_dir is not None else self.cfg.val_dir
        train_cache_file = os.path.splitext(train_dir)[0] + ".pkl"
        val_cache_file = os.path.splitext(val_dir)[0] + ".pkl"
        if self.cfg.sampling:
            train_dataset = base_dataset._build_dataset(data_dir=train_dir)
            train_dataloader = base_dataset.build_dataloader_from_dataset(
                train_dataset,
                batch_size=self.cfg.batch_size,
                processor=self.processor,
                device="cpu",
                shuffle=True,
                num_workers=self.cfg.num_workers,
                cache_file=train_cache_file,
                use_sampling=True,
            )
        else:
            train_dataloader = self._build_dataloader(
                train_dir if train_dir is not None else self.cfg.train_dir,
                cache_file=train_cache_file,
                use_sampling=False,
            )

        self.logger.info("Building valid dataloader...")
        val_dataloader = self._build_dataloader(
            val_dir if val_dir is not None else self.cfg.val_dir,
            cache_file= val_cache_file,
            use_sampling=False,
        )

        self.logger.info(
            f"Info dataset: train = {len(train_dataloader)}, test = {len(val_dataloader)}"
        )

        optimizer = AdamW(self.model.parameters(), lr=self.cfg.lr)
        if self.cfg.scheduler:
            # scheduler = torch.optim.lr_scheduler.OneCycleLR(
            #     optimizer, 
            #     max_lr=self.cfg.lr, 
            #     steps_per_epoch=len(train_dataloader), 
            #     epochs=self.cfg.epochs, 
            #     anneal_strategy='cos', 
            #     pct_start=0.1,
            #     div_factor=25,   #init lr = max_lr / div_factor
            #     final_div_factor=1e4, # min lr = init_lr / final_dev_factor
            # )
            num_training_steps = self.cfg.epochs * len(train_dataloader)
            scheduler = get_scheduler(
                name=self.cfg.scheduler, 
                optimizer=optimizer, 
                num_warmup_steps=0,
                num_training_steps=num_training_steps,

            )

        if self.cfg.wandb:
            wandb.config = dict(self.cfg)

        best_acc = 0.0
        best_epoch = 0
        yaml_save(os.path.join(self.cfg.save_dir, "config.yaml"), dict(self.cfg))
        for epoch in range(self.cfg.epochs):

            # sampling slice window
            if (
                self.cfg.sampling and epoch != 0 and epoch % self.cfg.sampling == 0
            ):  # sampling each cfg.sampling epochs
                train_dataloader = base_dataset.build_dataloader_from_dataset(
                    train_dataset,
                    batch_size=self.cfg.batch_size,
                    processor=self.processor,
                    device="cpu",
                    shuffle=True,
                    num_workers=self.cfg.num_workers,
                    cache_file=train_cache_file,
                    use_sampling=True,
                )

            self.model.train()
            self.logger.info(f"Epoch: {epoch}:")

            running_loss = 0.0
           
            for batch in tqdm(train_dataloader):
                # forward pass
                batch = self._to_device(batch, self.cfg.device)
                
                outputs = self.model(**batch)
                loss = outputs.loss
                running_loss += loss.item()
                # backward pass to get the gradients
                loss.backward()

                # update
                optimizer.step()
                if self.cfg.scheduler:
                    scheduler.step()
                optimizer.zero_grad()

            loss_avg = running_loss / len(train_dataloader)
            self.logger.info(f"Epoch[{epoch}/{self.cfg.epochs}] - lr: {round(scheduler.get_last_lr()[0], 9) if self.cfg.scheduler else self.cfg.lr} - loss: {loss_avg}")
            if self.cfg.wandb:
                wandb.log({"train_loss": loss_avg})

            # valid
            if epoch >= self.cfg.eval_delay:
                acc = self.val_on_dataloader(val_dataloader)

                if acc > best_acc:
                    self.model.save_pretrained(os.path.join(self.cfg.save_dir, "best"))
                    self.logger.info(f"Update best acc, prev best acc = {best_acc}, current best acc = {acc}")
                    best_acc = acc
                    best_epoch = epoch

                if epoch % self.cfg.save_weight_interval == 0:
                    self.model.save_pretrained(
                        os.path.join(self.cfg.save_dir, "epoch_{}".format(epoch))
                    )

            self.model.save_pretrained(os.path.join(self.cfg.save_dir, "last"))
        self.logger.info(f"Best accuracy = {best_acc} at epoch {best_epoch}")

    def _to_device(self, batch, device):
        batch = {k: v.to(device) for k, v in batch.items()}
        return batch

    def val_on_dataloader(self, dataloader):
        self.model.eval()
        total, correct = 0, 0
        preds, truths = [], []

        running_loss = 0.0
        for batch in tqdm(dataloader):
            with torch.no_grad():
                batch = self._to_device(batch, self.cfg.device)
                outputs = self.model(**batch)

                loss = outputs.loss
                running_loss += loss.item()
                predictions = outputs.logits.argmax(dim=2)
                valid_samples = batch["labels"] != -100
                predictions = predictions[valid_samples]
                batch_labels = batch["labels"][valid_samples]

                preds.extend(predictions.detach().cpu().numpy().tolist())
                truths.extend(batch_labels.detach().cpu().numpy().tolist())
                correct += (predictions == batch_labels).float().sum()
                total += predictions.numel()

        loss_avg = running_loss / len(dataloader)

        p, r, f1, support = precision_recall_fscore_support(truths, preds)
        # self.logger.info("shapeeee", p.shape)
        table_data = [["Class", "P", "R", "F1", "#samples"]]
        for c in range(len(self.cfg.classes)):
            if c < p.shape[0]:
                table_data.append([self.cfg.classes[c], p[c], r[c], f1[c], support[c]])
            continue

        f1_avg = sum(f1) / len(f1)
        table = AsciiTable(table_data)
        self.logger.info(table.table)
        self.logger.info(
            "Validation F1: {} - #samples: {} - #corrects: {}".format(
                f1_avg, total, correct
            )
        )

        if self.cfg.wandb:
            wandb.log({"val_loss": loss_avg, "f1": f1_avg})

        return f1_avg
