import os

# os.environ["CUDA_VISIBLE_DEVICES"] = "6"

import argparse

import json

from collections import defaultdict

import numpy as np

import pytorch_lightning as pl

import torch

import wandb

from pytorch_lightning.callbacks import ModelCheckpoint

from sklearn.metrics import accuracy_score, classification_report, jaccard_score, roc_auc_score

from torch.nn import BCEWithLogitsLoss

from torch.utils.data import DataLoader

from torchinfo import summary

from tqdm import tqdm

from transformers import AdamW

from findings_classifier.chexpert_dataset import Chexpert_Dataset

from findings_classifier.chexpert_model import ChexpertClassifier

from local_config import WANDB_ENTITY

class LitIGClassifier(pl.LightningModule):

    def __init__(self, num_classes, class_names, class_weights=None, learning_rate=1e-5):

        super().__init__()

        # Model

        self.model = ChexpertClassifier(num_classes)

        # Loss with class weights

        if class_weights is None:

            self.criterion = BCEWithLogitsLoss()

        else:

            self.criterion = BCEWithLogitsLoss(pos_weight=class_weights)

        # Learning rate

        self.learning_rate = learning_rate

        self.class_names = class_names

    def forward(self, x):

        return self.model(x)

    def step(self, batch, batch_idx):

        x, y = batch['image'].to(self.device), batch['labels'].to(self.device)

        logits = self(x)

        loss = self.criterion(logits, y)

        # Apply sigmoid to get probabilities

        preds_probs = torch.sigmoid(logits)

        # Get predictions as boolean values

        preds = preds_probs > 0.5

        # calculate jaccard index

        jaccard = jaccard_score(y.cpu().numpy(), preds.detach().cpu().numpy(), average='samples')

        class_report = classification_report(y.cpu().numpy(), preds.detach().cpu().numpy(), output_dict=True)

        # scores = class_report['micro avg']

        scores = class_report['macro avg']

        metrics_per_label = {label: metrics for label, metrics in class_report.items() if label.isdigit()}

        f1 = scores['f1-score']

        rec = scores['recall']

        prec = scores['precision']

        acc = accuracy_score(y.cpu().numpy().flatten(), preds.detach().cpu().numpy().flatten())

        try:

            auc = roc_auc_score(y.cpu().numpy().flatten(), preds_probs.detach().cpu().numpy().flatten())

        except Exception as e:

            auc = 0.

        return loss, acc, f1, rec, prec, jaccard, auc, metrics_per_label

    def training_step(self, batch, batch_idx):

        loss, acc, f1, rec, prec, jaccard, auc, _ = self.step(batch, batch_idx)

        train_stats = {'loss': loss, 'train_acc': acc, 'train_f1': f1, 'train_rec': rec, 'train_prec': prec, 'train_jaccard': jaccard,

                       'train_auc': auc}

        wandb_run.log(train_stats)

        return train_stats

    def training_epoch_end(self, outputs):

        avg_loss = torch.stack([x['loss'] for x in outputs]).mean()

        avg_acc = np.mean([x['train_acc'] for x in outputs])

        avg_f1 = np.mean([x['train_f1'] for x in outputs])

        avg_rec = np.mean([x['train_rec'] for x in outputs])

        avg_prec = np.mean([x['train_prec'] for x in outputs])

        avg_jaccard = np.mean([x['train_jaccard'] for x in outputs])

        avg_auc = np.mean([x['train_auc'] for x in outputs])

        wandb_run.log({'epoch_train_loss': avg_loss, 'epoch_train_acc': avg_acc, 'epoch_train_f1': avg_f1, 'epoch_train_rec': avg_rec,

                       'epoch_train_prec': avg_prec, 'epoch_train_jaccard': avg_jaccard, 'epoch_train_auc': avg_auc})

    def validation_step(self, batch, batch_idx):

        loss, acc, f1, rec, prec, jaccard, auc, metrics_per_label = self.step(batch, batch_idx)

        # log f1 for checkpoint callback

        self.log('val_f1', f1)

        return {'val_loss': loss, 'val_acc': acc, 'val_f1': f1, 'val_rec': rec, 'val_prec': prec, 'val_jaccard': jaccard,

                'val_auc': auc}, metrics_per_label

    def validation_epoch_end(self, outputs):

        outputs, per_label_metrics_outputs = zip(*outputs)

        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()

        avg_acc = np.mean([x['val_acc'] for x in outputs])

        avg_f1 = np.mean([x['val_f1'] for x in outputs])

        avg_rec = np.mean([x['val_rec'] for x in outputs])

        avg_prec = np.mean([x['val_prec'] for x in outputs])

        avg_jaccard = np.mean([x['val_jaccard'] for x in outputs])

        avg_auc = np.mean([x['val_auc'] for x in outputs])

        per_label_metrics = defaultdict(lambda: defaultdict(float))

        label_counts = defaultdict(int)

        for metrics_per_label in per_label_metrics_outputs:

            for label, metrics in metrics_per_label.items():

                label_name = self.class_names[int(label)]

                per_label_metrics[label_name]['precision'] += metrics['precision']

                per_label_metrics[label_name]['recall'] += metrics['recall']

                per_label_metrics[label_name]['f1-score'] += metrics['f1-score']

                per_label_metrics[label_name]['support'] += metrics['support']

                label_counts[label_name] += 1

        # Average the metrics

        for label, metrics in per_label_metrics.items():

            for metric_name in ['precision', 'recall', 'f1-score']:

                if metrics['support'] > 0:

                    per_label_metrics[label][metric_name] /= label_counts[label]

        val_stats = {'val_loss': avg_loss, 'val_acc': avg_acc, 'val_f1': avg_f1, 'val_rec': avg_rec, 'val_prec': avg_prec, 'val_jaccard': avg_jaccard,

                     'val_auc': avg_auc}

        wandb_run.log(val_stats)

    def test_step(self, batch, batch_idx):

        loss, acc, f1, rec, prec, jaccard, auc, _ = self.step(batch, batch_idx)

        return {'test_loss': loss, 'test_acc': acc, 'test_f1': f1, 'test_rec': rec, 'test_prec': prec, 'test_jaccard': jaccard, 'test_auc': auc}

    def test_epoch_end(self, outputs):

        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()

        avg_acc = np.mean([x['test_acc'] for x in outputs])

        avg_f1 = np.mean([x['test_f1'] for x in outputs])

        avg_rec = np.mean([x['test_rec'] for x in outputs])

        avg_prec = np.mean([x['test_prec'] for x in outputs])

        avg_jaccard = np.mean([x['test_jaccard'] for x in outputs])

        avg_auc = np.mean([x['test_auc'] for x in outputs])

        test_stats = {'test_loss': avg_loss, 'test_acc': avg_acc, 'test_f1': avg_f1, 'test_rec': avg_rec, 'test_prec': avg_prec,

                      'test_jaccard': avg_jaccard, 'test_auc': avg_auc}

        wandb_run.log(test_stats)

    def configure_optimizers(self):

        optimizer = AdamW(self.parameters(), lr=self.learning_rate)

        return optimizer

def save_preds(dataloader, split):

    # load checkpoint

    ckpt_path = f"findings_classifier/checkpoints/chexpert_train/ChexpertClassifier-epoch=06-val_f1=0.36.ckpt"

    model = LitIGClassifier.load_from_checkpoint(ckpt_path, num_classes=num_classes, class_weights=val_dataset.get_class_weights(),

                                                 class_names=class_names, learning_rate=args.lr)

    model.eval()

    model.cuda()

    model.half()

    class_names_np = np.asarray(class_names)

    # get predictions for all study ids

    structured_preds = {}

    for batch in tqdm(dataloader):

        dicom_ids = batch['dicom_id']

        logits = model(batch['image'].half().cuda())

        preds_probs = torch.sigmoid(logits)

        preds = preds_probs > 0.5

        # iterate over each study id in the batch

        for i, (dicom_id, pred) in enumerate(zip(dicom_ids, preds.detach().cpu())):

            # get all positive labels

            findings = class_names_np[pred].tolist()

            structured_preds[dicom_id] = findings

    # save predictions

    with open(f"findings_classifier/predictions/structured_preds_chexpert_log_weighting_macro_{split}.json", "w") as f:

        json.dump(structured_preds, f, indent=4)

if __name__ == '__main__':

    parser = argparse.ArgumentParser()

    parser.add_argument("--run_name", type=str, default="debug")

    parser.add_argument("--lr", type=float, default=5e-5)

    parser.add_argument("--epochs", type=int, default=6)

    parser.add_argument("--loss_weighting", type=str, default="log", choices=["lin", "log", "none"])

    parser.add_argument("--truncate", type=int, default=None)

    parser.add_argument("--batch_size", type=int, default=64)

    parser.add_argument("--num_workers", type=int, default=12)

    parser.add_argument("--use_augs", action="store_true", default=False)

    parser.add_argument("--train", action="store_true", default=False)

    args = parser.parse_args()

    TRAIN = args.train

    # fix all seeds

    pl.seed_everything(42, workers=True)

    # Create DataLoaders

    train_dataset = Chexpert_Dataset(split='train', truncate=args.truncate, loss_weighting=args.loss_weighting, use_augs=args.use_augs)

    val_dataset = Chexpert_Dataset(split='validate', truncate=args.truncate)

    test_dataset = Chexpert_Dataset(split='test')

    train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers)

    val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, num_workers=args.num_workers)

    test_dataloader = DataLoader(test_dataset, batch_size=args.batch_size, num_workers=args.num_workers)

    # Number of classes for IGClassifier

    num_classes = len(train_dataset.chexpert_cols)

    class_names = train_dataset.chexpert_cols

    if TRAIN:

        class_weights = torch.tensor(train_dataset.get_class_weights(), dtype=torch.float32)

        # Define the model

        lit_model = LitIGClassifier(num_classes, class_names=class_names, class_weights=class_weights, learning_rate=args.lr)

        print(summary(lit_model))

        # WandB logger

        wandb_run = wandb.init(

            project="ChexpertClassifier",

            entity= WANDB_ENTITY,

            name=args.run_name

        )

        # checkpoint callback

        checkpoint_callback = ModelCheckpoint(

            monitor='val_f1',

            dirpath=f'findings_classifier/checkpoints/{args.run_name}',

            filename='ChexpertClassifier-{epoch:02d}-{val_f1:.2f}',

            save_top_k=1,

            save_last=True,

            mode='max',

        )

        # Train the model

        trainer = pl.Trainer(max_epochs=args.epochs, gpus=1, callbacks=[checkpoint_callback], benchmark=False, deterministic=True, precision=16)

        trainer.fit(lit_model, train_dataloader, val_dataloader)

        # Test the model

        # trainer.validate(lit_model, val_dataloader, ckpt_path="checkpoints_IGCLassifier/lr_5e-5_to0_log_weighting_patches_augs_imgemb/IGClassifier-epoch=09-val_f1=0.65.ckpt")

    else:

        save_preds(train_dataloader, "train")

        save_preds(val_dataloader, "val")

        save_preds(test_dataloader, "test")