from pathlib import Path

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader, TensorDataset
from torchvision import models, transforms
from torchvision.transforms import v2, GaussianBlur
from tqdm import tqdm
import numpy as np
import time
from efficientnet_pytorch import EfficientNet
from nets_parts.datasets.psi_dataset_regions_cls import (
    PSIRegionClsDataset,
    PSIRegionClsDatasetParams,
)
from nets_parts.datasets.psi_dataset_random import (
    PSIRandomDataset,
    PSIRandomDatasetParams,
)
from nets_parts.nets_train_part import (
    run_valid, run_train,
    save_confusion_matrix,
    get_data_iterator, get_data_iterator_pl
)
from torchvision.transforms import RandAugment, ToPILImage
from torchvision.transforms.functional import pil_to_tensor
import gc
from nets_parts.datasets.psi_torch_dataset import TorchPSIDataset
from nets_parts.datasets.utils import low_entropy_filter
from torch.utils.tensorboard import SummaryWriter
import sys, json, os
from os.path import isdir, join
from nets_parts.datasets.utils import low_entropy_filter_single_input_with_prob
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
print(os.path.dirname(os.path.realpath(__file__)))
from nets_parts.RandAugment import RandAugment as MyRandAugment
THRESHOLD = 0.95

def train_model(
    model,
    criterion,
    optimizer,
    scheduler,
    iterations_per_epoch,
    n_epochs,
    valid_loader,
    patch_size,
    layer,
    batch_size_train,
    valid_image_num,
    checkpoint_path,
    cur_statistic_path,
    rand_aug_strange,
    rand_aug_max_strange,
    n_ops,
    epoch_for_start_pl,
    use_all_data,
    writer,
    use_my_rand_aug,
    images_to_use
):
    global_train_loss = []
    global_train_acc = []
    global_valid_loss = []
    global_valid_acc = []
    best_train_acc = None
    best_valid_acc = None
    train_data_pseudo_label = None
    train_ds = None

    data_pl_iterator, train_data_pseudo_label = get_data_iterator_pl(
        path_to_data="/home/n.yakovlev/datasets/symblink/WSS2/train_valid",
        train_ds=train_data_pseudo_label,
        layer=layer,
        patch_size=patch_size,
        batch_size=64
    )
    generated_images: int = 0
    if use_my_rand_aug:
        RandAugmentator = MyRandAugment(ops_num=n_ops, cur_value=rand_aug_strange, max_value=rand_aug_max_strange)
    else:
        RandAugmentator = RandAugment(num_ops=n_ops, magnitude=rand_aug_strange, num_magnitude_bins=rand_aug_max_strange, fill=255)
    using_labels_in_cur_epoch = []
    for epoch in range(n_epochs):
        model.train()
        if epoch % images_to_use == 0:
            data_iterator, train_ds = get_data_iterator(
                path_to_data=f"/home/n.yakovlev/datasets/symblink/WSS2/train_valid_{images_to_use}",
                train_ds=train_ds,
                layer=layer,
                patch_size=patch_size,
                batch_size=64
            )
        print(f"Epoch {epoch}/{n_epochs - 1}")
        print("-" * 10)

        # Each epoch has a training and validation phase
        model.train()  # Set model to training mode
        torch.set_grad_enabled(True)

        running_loss = 0.0
        running_corrects = 0
        using_labels_in_cur_epoch.append({k: int(0) for k in range(5)})

        # Iterate over data.
        for _ in tqdm(
            range(iterations_per_epoch), f"running epoch {epoch + 1}"
        ):
            model.eval()
            torch.set_grad_enabled(False)
            inputs_pl = None
            labels_pl = None
            cycle_try_to_get_pl = 0
            while epoch >= epoch_for_start_pl and (inputs_pl is None or inputs_pl.size(dim=0) < batch_size_train * 7):
                cycle_try_to_get_pl += 1
                gc.collect()
                torch.cuda.empty_cache()
                if inputs_pl is None and cycle_try_to_get_pl > 10:
                    break
                if generated_images > train_data_pseudo_label.__len__() - 10 * batch_size_train:
                    generated_images = 0
                    data_pl_iterator, train_data_pseudo_label = get_data_iterator_pl(
                        path_to_data="/home/n.yakovlev/datasets/symblink/WSS2/train_valid",
                        train_ds=train_data_pseudo_label,
                        layer=layer,
                        patch_size=patch_size,
                        batch_size=64
                    )

                images_cur = next(data_pl_iterator).to(device)
                generated_images += images_cur.size(dim=0)
                labels_cur = model(images_cur)
                if (labels_cur.max(dim=1).values > THRESHOLD).sum().item() < 10:
                    continue
                # get pseudo labels with threshold
                images_cur = images_cur[labels_cur.max(dim=1).values > THRESHOLD]
                labels_cur = labels_cur[labels_cur.max(dim=1).values > THRESHOLD]
                labels_cur[labels_cur > THRESHOLD] = 1.0
                labels_cur[labels_cur < 1 - 1e-5] = 0
                if inputs_pl is None:
                    inputs_pl = images_cur
                    labels_pl = labels_cur
                else:
                    inputs_pl = torch.cat((inputs_pl, images_cur), dim=0)
                    labels_pl = torch.cat((labels_pl, labels_cur), dim=0)
            model.train()  # Set model to training mode
            torch.set_grad_enabled(True)
            if not (inputs_pl is None or labels_pl is None):
                inputs_pl, labels_pl = inputs_pl[:7 * batch_size_train], labels_pl[:7 * batch_size_train]
            if inputs_pl is None or labels_pl is None or inputs_pl.size(0) == 0:
                inputs, labels = next(data_iterator)
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                outputs = model(inputs)
                loss = criterion(outputs, labels)

                loss.backward()
                optimizer.step()

                # statistics
                preds = torch.argmax(outputs, dim=1)
                gts = torch.argmax(labels, dim=1)
                running_loss += loss.item()
                running_corrects += torch.sum(preds == gts)
            else:
                uniq_values, uniq_counts = np.unique(labels_pl.argmax(1).cpu().numpy(), return_counts=True)
                for ind, i in zip(uniq_values, uniq_counts):
                    using_labels_in_cur_epoch[epoch][ind] += int(i)
                for image in inputs_pl:
                    if use_my_rand_aug:
                        temp_img = (image * 255).to(torch.uint8)
                        temp_img = ToPILImage()(temp_img)
                        temp_img = RandAugmentator(temp_img)
                        image = pil_to_tensor(temp_img) / 255
                    else:
                        image = (RandAugmentator((image * 255).to(torch.uint8))).to(torch.float32)
                        if image.max() > 1.1:
                            image /= 255
                # Each epoch has a training and validation phase
                inputs, labels = next(data_iterator)
                labels_pl = labels_pl.to(labels.dtype)
                inputs = inputs.to(device)
                labels = labels.to(device)
                inputs_pl = inputs_pl.to(device)
                labels_pl = labels_pl.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                outputs = model(inputs)
                loss_l = criterion(outputs, labels)

                outputs_pl = model(inputs_pl)
                loss_u = criterion(outputs_pl, labels_pl)

                loss = loss_l + 1. / 7 * loss_u
                loss.backward()
                optimizer.step()

                # statistics
                preds = torch.argmax(outputs, dim=1)
                gts = torch.argmax(labels, dim=1)
                running_loss += loss.item()
                running_corrects += torch.sum(preds == gts)
            gc.collect()
            torch.cuda.empty_cache()

        if epoch >= 10:
            scheduler.step()

        epoch_loss = running_loss / iterations_per_epoch
        epoch_acc = running_corrects.float() / (
            iterations_per_epoch * inputs.size(0)
        )
        global_train_acc.append(float(epoch_acc))
        global_train_loss.append(float(epoch_loss))
        if best_train_acc is None or best_train_acc < epoch_acc:
            torch.save(model.state_dict(), f"{checkpoint_path}/best_train_acc_{round(float(epoch_acc) * 100)}.pth")
            best_train_acc = epoch_acc

        print(f"Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}")
        writer.add_scalar("Loss", epoch_loss, epoch)
        writer.add_scalar("Acc", epoch_acc, epoch)

        matrix_conf, epoch_loss, epoch_acc = run_valid(
            model=model,
            valid_loader=valid_loader,
            valid_images_num=valid_image_num,
            criterion=criterion,
            device=device
        )
        global_valid_acc.append(float(epoch_acc))
        global_valid_loss.append(float(epoch_loss))
        if best_valid_acc is None or best_valid_acc < epoch_acc:
            torch.save(model.state_dict(), f"{checkpoint_path}/best_valid_acc_{round(float(epoch_acc) * 100)}.pth")
            best_valid_acc = epoch_acc
        print(f"Val Loss: {epoch_loss:.4f} Val Acc: {epoch_acc:.4f}")
        writer.add_scalar("Loss valid", epoch_loss, epoch)
        writer.add_scalar("Accuracy valid", epoch_acc, epoch)
        if epoch % 1 == 0:
            with open(f"{cur_statistic_path}/statistic.json", "w") as f:
                json.dump(
                    {
                        "train_loss": global_train_loss,
                        "train_acc": global_train_acc,
                        "valid_loss": global_valid_loss,
                        "valid_acc": global_valid_acc,
                    },
                    f
                )
            with open(f"{cur_statistic_path}/using_pl.json", "w") as f:
                json.dump(
                    using_labels_in_cur_epoch,
                    f
                )
            path_to_save_conf_matrix = f"{cur_statistic_path}/confusion_matrix_best_epoch.png"
            save_confusion_matrix(matrix_conf, path_to_save_conf_matrix)
    if not (train_data_pseudo_label is None):
        train_data_pseudo_label.close()
    if not (train_ds is None):
        train_ds.close()
    train_data_pseudo_label = None
    train_ds = None


if __name__ == "__main__":
    """
    {
        "batch_size_train": 192,
        "batch_size_valid": 512,
        "lr": 0.01,
        "exprement_path": "/home/n.yakovlev/my_best_program/diplom_8sem/experiments/classifier",

        "patch_size": 224,
        "layer": 2,
        "scheduler_param": 0.99,
        "pretrain": 1,

        "device": 1,

        "is_efficientnet": 1
        "nn_name": "efficientnet-b4"
        "iter_per_epoch": 26
    }
"""
    path_to_params = sys.argv[1]
    with open(path_to_params) as f:
        parsed_file = f.read()
        parsed_json = json.loads(parsed_file)
        BATCH_SIZE_TRAIN = int(parsed_json["batch_size_train"])
        BATCH_SIZE_VALID = int(parsed_json["batch_size_valid"])
        LEARING_RATE = float(parsed_json["lr"])
        LAYER = int(parsed_json["layer"])

        PRETRAIN = bool(parsed_json["pretrain"])
        EXP_PATH = str(parsed_json["exprement_path"])
        os.makedirs(EXP_PATH, exist_ok=True)
        PATCH_SIZE = int(parsed_json["patch_size"])
        SCHEDULER_PARAM = float(parsed_json["scheduler_param"])

        IS_EFFICIENTNET = bool(parsed_json["is_efficientnet"])
        NN_NAME = str(parsed_json["nn_name"])
        ITER_PER_EPOCH = int(parsed_json["iter_per_epoch"])
        rand_aug_strange = int(parsed_json["rand_aug_strange"])
        rand_aug_max_strange = int(parsed_json["rand_aug_max_strange"])
        n_ops = int(parsed_json["n_ops"])
        device = f"cuda:{parsed_json["device"]}"
        epoch_for_start_pl = int(parsed_json["epoch_for_start_pl"])
        use_all_data = bool(parsed_json["use_all_data"])
        use_my_rand_aug = bool(parsed_json["use_my_rand_aug"])
        images_to_use = int(parsed_json["images_to_use"])

    data = np.load("/home/n.yakovlev/datasets/test_files_WSS2.npz")
    images = data["images"]
    valid_image_num = images.shape[0]
    labels = data["labels"]

    del data
    gc.collect()
    valid_dataset = TensorDataset(torch.tensor(images), torch.tensor(labels))
    del images, labels
    valid_loader = DataLoader(valid_dataset, batch_size=BATCH_SIZE_VALID, num_workers=8)

    for use_my_rand_aug in [False]:
        for rand_aug_strange in [4, 5, 6, 7, 8, 9, 10, 11, 12, 13]:
            rand_aug_max_strange = 31
            n_ops = int(parsed_json["n_ops"])
            path_to_params = sys.argv[1]
            with open(path_to_params) as f:
                parsed_file = f.read()
                parsed_json = json.loads(parsed_file)
                parsed_json["rand_aug_strange"] = rand_aug_strange
                parsed_json["use_my_rand_aug"] = use_my_rand_aug

                EXP_PATH = str(parsed_json["exprement_path"])
                os.makedirs(EXP_PATH, exist_ok=True)
                onlydirs = [f for f in os.listdir(EXP_PATH) if isdir(join(EXP_PATH, f))]
                cur_exp_path = f"{EXP_PATH}/{len(onlydirs)}"
                os.makedirs(cur_exp_path)
                cur_statistic_path = f"{cur_exp_path}/statistic"
                os.makedirs(cur_statistic_path)
                checkpoint_path = cur_exp_path + "/checkpoints"
                os.makedirs(checkpoint_path)
                writer = SummaryWriter(EXP_PATH)
                with open(f"{cur_exp_path}/params.json", 'w') as f:
                    json.dump(parsed_json, f)

            if IS_EFFICIENTNET:
                if PRETRAIN:
                    model_ft = EfficientNet.from_pretrained(NN_NAME, num_classes=5)
                else:
                    model_ft = EfficientNet.from_name(NN_NAME, num_classes=5)
            else:
                model_ft = models.resnet50(weights="IMAGENET1K_V1")
                num_ftrs = model_ft.fc.in_features
                model_ft.fc = nn.Linear(num_ftrs, 5)

            try:
                pretrained_dict = torch.load("/home/n.yakovlev/conv_autoencoder_256neuron_best_valid.pth")
                model_dict = model_ft.state_dict()

                processed_dict = {}

                for k in model_dict.keys():
                    decomposed_key = k.split(".")
                    if ("model" in decomposed_key):
                        pretrained_key = ".".join(decomposed_key[1:])
                        processed_dict[k] = pretrained_dict[pretrained_key]

                model_ft.load_state_dict(processed_dict, strict=False)

            except Exception:
                print("Model weights doesn't loaded!")


            model_ft = model_ft.to(device)

            criterion = nn.CrossEntropyLoss(weight=torch.tensor([1, 1, 1.2, 1, 1.2], device=device))

            # Observe that all parameters are being optimized
            optimizer_ft = optim.Adam(model_ft.parameters(), lr=LEARING_RATE)

            exp_lr_scheduler = lr_scheduler.StepLR(
                optimizer_ft, step_size=1, gamma=SCHEDULER_PARAM
            )

            train_model(
                model_ft,
                criterion,
                optimizer_ft,
                exp_lr_scheduler,
                iterations_per_epoch=ITER_PER_EPOCH,  # len(train_ds) // batch_size,
                n_epochs=50,
                valid_loader=valid_loader,
                patch_size=PATCH_SIZE,
                layer=LAYER,
                batch_size_train=BATCH_SIZE_TRAIN,
                valid_image_num=valid_image_num,
                checkpoint_path=checkpoint_path,
                cur_statistic_path=cur_statistic_path,
                rand_aug_strange=rand_aug_strange,
                rand_aug_max_strange=rand_aug_max_strange,
                n_ops=n_ops,
                epoch_for_start_pl=epoch_for_start_pl,
                use_all_data=use_all_data,
                writer=writer,
                use_my_rand_aug=use_my_rand_aug,
                images_to_use=images_to_use
            )
            del model_ft, optimizer_ft
            gc.collect()
            torch.cuda.empty_cache()