from pathlib import Path

import numpy as np
import torch
import yaml
from nnfabrik import builder
from tqdm import tqdm

import wandb

from models.ln_model import FactorizedModel, RankModel

from training.measures import correlation
from training.train_steps import train_step, save_checkpoint, model_step
from utils.global_functions import model_seed


def train_ln(
        dataset_fn,
        dataloader_config,
        model_config,
        dataset_config_dict,
        model_dir,
        device,
        max_coordinate,
        img_h,
        img_w,
        wandb_log,
        config,
        filter_type,
        cuda,
        performance,
        nm=False,
        mei=None,
        save_every_epoch=False,
        flip_sta=False,
        cs=False
):
    dataloaders = builder.get_data(dataset_fn, dataloader_config)
    print(dataloaders)

    first_session_ID = list((dataloaders["train"].keys()))[0]
    print(first_session_ID)
    a_dataloader = dataloaders["train"][first_session_ID]

    img_h = int(
        (img_h - a_dataloader.dataset.crop[0] - a_dataloader.dataset.crop[1])
        / dataloader_config["subsample"]
    )
    img_w = int(
        (img_w - a_dataloader.dataset.crop[2] - a_dataloader.dataset.crop[3])
        / dataloader_config["subsample"]
    )


    sta = model_config['sta']

    if not cs:
        model_config["input_shape"] = sta[0].shape
        model = FactorizedModel(**model_config)
        basic = ''
    else:
        model_config["input_shape"] = sta[0].shape
        model = RankModel(**model_config)
        basic = 'cs-'

    for name, parameter in model.named_parameters():
        if 'bias' in name:
            parameter.data = torch.zeros(parameter.shape)
            parameter.requires_grad = False
    if cuda:
        model = model.to(device)

    model = model.double()

    loss_function = config["loss_function"](**config["loss_function_config"])

    optimizer = config["optimizer"](model.parameters(), **config["optimizer_config"])

    scheduler = config["scheduler"](optimizer=optimizer, **config["scheduler_config"])

    weight_dir = f'{config["log_dir"]}/{config["model_dir"]}/weights/seed_{model_seed}/'
    stats_dir = f'{config["log_dir"]}/{config["model_dir"]}/stats/seed_{model_seed}/'
    config_dir = f'{config["log_dir"]}/{config["model_dir"]}/config/'

    Path(weight_dir).mkdir(exist_ok=True, parents=True)
    Path(stats_dir).mkdir(exist_ok=True, parents=True)
    Path(config_dir).mkdir(exist_ok=True, parents=True)

    max_avg_valid_corr = -1
    nm_str = 'nm' if nm else 'wn'

    with open(f"{config_dir}/config.yaml", "w") as file:
        yaml.dump(config, file)
    if wandb_log:
        wandb.init(
            config=config,
            project=f"{basic}{filter_type}-model-{dataset_config_dict['data_type']}_{nm_str}",
            entity="retinal-circuit-modeling", dir=config['log_dir']
        )
        wandb.define_metric('epoch')
        wandb.define_metric('batch')

    score, correlations = train(
        epochs=config["epochs"],
        dataloaders=dataloaders,
        dataset_config=dataloader_config,
        model=model,
        optimizer=optimizer,
        loss_function=loss_function,
        scheduler=scheduler,
        device=device,
        model_dir=model_dir,
        weight_dir=weight_dir,
        stats_dir=stats_dir,
        stopper_patience=config["stopper_patience"],
        regularizer_start=config["regularizer_start"],
        size=model_config['input_shape'],
        max_coordinate=max_coordinate,
        img_h=img_h,
        img_w=img_w,
        config_dict=dataset_config_dict,
        wandb_log=wandb_log,
        fancy_nonlin=model_config["fancy_nonlin"],
        learn_filter=model_config["learn_filter"],
        performance=performance,
        save_every_epoch=save_every_epoch
    )
    if wandb_log:
        wandb.finish()
    print(score, correlations)

def train(
        epochs,
        dataloaders,
        dataset_config,
        config_dict,
        model,
        optimizer,
        loss_function,
        scheduler,
        device,
        model_dir,
        weight_dir,
        stats_dir,
        stopper_patience,
        size,
        max_coordinate=None,
        img_h=150,
        img_w=200,
        regularizer_start=0,
        wandb_log=True,
        fancy_nonlin=False,
        learn_filter=True,
        start_epoch=0,
        continued_training=False,
        performance="validation",
        multi_nonlin=False,
        save_every_epoch=False,
):

    correlations = []
    train_correlations = []
    valid_losses = []
    train_losses = []
    penalties = []

    patience = stopper_patience
    max_avg_valid_corr = -1
    if not learn_filter and not fancy_nonlin:
        epochs = 1
    if performance == "validation":
        perf_str = "valid"
    else:
        perf_str = performance

    save_checkpoint(0, model, path=f"{weight_dir}/initial_model.m",
                    optimizer=optimizer,
                    valid_corr='-1')
    if isinstance(dataset_config["retina_index"], int):
        dataset_config['retina_index'] = str(dataset_config["retina_index"] + 1).zfill(2)
    for epoch in range(start_epoch, epochs + start_epoch):
        print(f"Epoch {epoch}")
        model.train()

        epoch_train_correlations = []
        epoch_valid_correlations = []
        epoch_train_losses = []
        epoch_penalties = []
        epoch_valid_losses = []

        print(size)
        rf_size = size

        if performance != "test":
            all_predictions = []
            all_responses = []
        if learn_filter or fancy_nonlin:
            for i, (images, responses) in enumerate(
                    tqdm(
                        dataloaders["train"][dataset_config["retina_index"]]
                    )
            ):

                responses = responses.transpose(1, 2).to(device).contiguous()
                responses = torch.flatten(responses, start_dim=0, end_dim=1)
                images = images.to(device)

                loss, penalty,  corr = train_step(
                    images=images.to(device).contiguous(),
                    responses=responses.to(device).contiguous(),
                    model=model,
                    loss_function=loss_function,
                    optimizer=optimizer,
                    rf_size=rf_size,
                    max_coordinate=max_coordinate,
                    h=img_h,
                    w=img_w,
                    penalty_coef=1,
                    regularizer_start=regularizer_start,
                    epoch=epoch,
                    wandb_log=wandb_log
                )
                epoch_train_losses.append(float(loss.item()))
                epoch_penalties.append([x.item() for x in penalty])
                epoch_train_correlations.append(np.mean(corr.detach().cpu().numpy()[0]))

            epoch_penalties = np.array(epoch_penalties)
            train_losses.append(sum(epoch_train_losses) / len(epoch_train_losses))
            penalties.append(np.mean(epoch_penalties, axis=0))
            train_corr = np.mean(epoch_train_correlations, axis=0)
            train_correlations.append(train_corr)
            print("")
            print(model_dir)
            print("avg train correlation", train_correlations[-1])
            print("avg train loss: ", train_losses[-1])
            print(
                "spatial penalty:",
                penalties[-1][0],
                "temporal_penalty:",
                penalties[-1][1],
                "l1 penalty:",
                penalties[-1][2],
            )
            if wandb_log:
                wandb.log(
                    {
                        "loss/train": train_losses[-1],
                        "spatial_penalty/train": penalties[-1][0],
                        "temporal_penalty/train": penalties[-1][1],
                        "l1_penalty/train": penalties[-1][2],
                        "correlation/train": train_corr,
                        "epoch": epoch
                    },

                )
                print('model:', model)
                if hasattr(model, 'conv1.mean_x'):
                    wandb.log(
                        {
                            'mean_x': model.conv1.mean_x,
                            'mean_y': model.conv1.mean_y,
                            'sigma_x': model.conv1.sigma_x,
                            'sigma_y': model.conv1.sigma_y,
                            "epoch": epoch
                        },

                    )
                if multi_nonlin:
                    wandb.log(
                        {
                            "a": model.nonlinearity.a[dataset_config["retina_index"]].detach().cpu(),
                            "b": model.nonlinearity.b[dataset_config["retina_index"]].detach().cpu(),
                            "w": model.nonlinearity.w[dataset_config["retina_index"]].detach().cpu(),
                            "epoch": epoch
                        },
                    )
                elif fancy_nonlin:
                    wandb.log(
                        {
                            "a": model.nonlinearity.a.detach().cpu().item(),
                            "b": model.nonlinearity.b.detach().cpu().item(),
                            "w": model.nonlinearity.w.detach().cpu().item(),
                            "epoch": epoch
                        }
                    )

        model.eval()

        with torch.no_grad():
            k = 0
            for images, responses in tqdm(
                    dataloaders[performance][
                        str(dataset_config["retina_index"])
                    ]
            ):
                images = images.double().to(device).contiguous()
                responses = responses.transpose(1, 2)
                responses = (
                    torch.flatten(responses, start_dim=0, end_dim=1)
                    .to(device)
                    .contiguous()
                )
                output = model_step(
                    images=images,
                    model=model,
                    max_coordinate=max_coordinate,
                    rf_size=rf_size,
                    h=img_h,
                    w=img_w,
                )
                corr = correlation(output, responses, 1e-12)
                valid_loss = loss_function(output, responses)
                k += 1

                if device != "cpu":
                    epoch_valid_correlations.append(
                        np.mean(corr.detach().cpu().numpy()[0])
                    )
                    epoch_valid_losses.append(float(valid_loss.detach().cpu().numpy()))
                else:
                    epoch_valid_correlations.append(np.mean(corr.detach().numpy()[0]))
                    epoch_valid_losses.append(float(valid_loss.detach().numpy()))
                if performance != "train":
                    all_predictions.append(output)
                    all_responses.append(responses)

            print(model_dir)
            print(f"max batch {perf_str} corr:", max(epoch_valid_correlations))
            single_correlations = np.mean(epoch_valid_correlations, axis=0)
            if performance != "train":
                print('getting all predictions performance')
                all_outputs = torch.cat(all_predictions)
                all_responses = torch.cat(all_responses)
                all_single_correlations = correlation(
                    all_outputs, all_responses, 1e-12
                ).mean().item()
                print(f"all avg {perf_str} corr:", all_single_correlations)

            print(f"avg {perf_str} corr:", single_correlations)
            print(
                f"avg {perf_str} loss:",
                sum(epoch_valid_losses) / len(epoch_valid_losses),
            )
            print(f"max {perf_str} loss:", max(epoch_valid_losses))
            if wandb_log:
                wandb.log(
                    {
                        f"correlation/{perf_str}": single_correlations,
                        "loss/valid": sum(epoch_valid_losses) / len(epoch_valid_losses),
                        "epoch": epoch
                    },

                )
            if save_every_epoch:
                print(
                    f"Saving model epch {epoch}:{single_correlations}"
                )
                save_checkpoint(
                    epoch,
                    model=model,
                    path=f"{weight_dir}/model_e_{epoch}.m",
                    optimizer=optimizer,
                    valid_corr=single_correlations,
                )
            if single_correlations > max_avg_valid_corr:
                patience = stopper_patience
                print(
                    f"Saving so far best model:{single_correlations}, previous best:{max_avg_valid_corr}"
                )
                max_avg_valid_corr = single_correlations
                save_checkpoint(
                    epoch,
                    model=model,
                    path=f"{weight_dir}/best_model.m",
                    optimizer=optimizer,
                    valid_corr=single_correlations,
                )
            else:
                patience -= 1
                if patience == 0:
                    break
            print(f'patience: {patience}', f'stopper_patiens: {stopper_patience}')

            valid_losses.append(sum(epoch_valid_losses) / len(epoch_valid_losses))
            correlations.append(single_correlations)
            scheduler.step(single_correlations)

        with torch.no_grad():
            np_correlations = np.array(correlations)
            np_train_correlations = np.array(train_correlations)
            np_train_losses = np.array(train_losses)
            np_valid_losses = np.array(valid_losses)
            if performance == "validation":
                np.save(f"{stats_dir}/correlations", np_correlations)
            else:
                np.save(f"{stats_dir}/correlations_{perf_str}", np_correlations)
            np.save(f"{stats_dir}/train_correlations", np_train_correlations)
            np.save(f"{stats_dir}/train_losses", np_train_losses)
            np.save(f"{stats_dir}/{perf_str}_losses", np_valid_losses)
    if wandb_log:
        wandb.finish()
    return np.max(np_correlations), correlations
