from pathlib import Path

from torch import optim

from datasets.stas import get_cell_sta, create_cell_file_from_config_version
from evaluations.ln_model_performance import get_model_and_dataloader_for_ln_based_on_setting
from training.trainers import train, train_ln
import wandb


def train_ln_cross_stimulus(model_dir,
                         filename,
                         directory_prefix,
                         data_dir,
                         seed,
                         config_dict,
                         retina,
                         cell_index,
                         loss_function=None,
                         scheduler=None,
                         optimizer=None,
                         fixation_file=None,
                         stimulus_seed=None,
                         lr=None,
                         data_type='marmoset',
                         setting='wn_for_nm',
                         device='cuda',
                         freeze_filter=False,
                         change_nonlin=True,
                         stopper_patience=50,
                         wandb_log=True,
                         log_dir = None,
                         flip_sta=None,
                         save_every_epoch=False,
                         flip_imgs=False,
                         fancy_nonlin=True,
                        learn_filter=True,
                        rank_model=False):
    base_path = f'models/ln_models/{directory_prefix}/'
    dataloaders, model, config = get_model_and_dataloader_for_ln_based_on_setting(
        setting=setting,
        filename=filename,
        cell_number=cell_index,
        retina_index=retina,
        device=device,
        data_dir=data_dir,
        model_dir=model_dir,
        seed=seed,
        data_type=data_type,
        fixation_file=fixation_file,
        stimulus_seed=stimulus_seed,
        freeze=False,
        flip_imgs=flip_imgs,
        fancy_nonlin=fancy_nonlin,
        config=config_dict,
        base_path=f'models/ln_models/{directory_prefix}/',
        factorized=True,
        rank_model=rank_model
    )
    weight_dir = f'{model_dir}/{base_path}/{data_type}/retina{retina}/cell_{cell_index}/{filename}/{setting}/ff_{1 if freeze_filter else 0}_pn_{1 if change_nonlin else 0}/weights/seed_{seed}/'
    stats_dir = f'{model_dir}/{base_path}/{data_type}/retina{retina}/cell_{cell_index}/{filename}/{setting}/ff_{1 if freeze_filter else 0}_pn_{1 if change_nonlin else 0}/stats/seed_{seed}/'
    model_dir = f'{model_dir}/{base_path}/{data_type}/retina{retina}/cell_{cell_index}/{filename}/{setting}/ff_{1 if freeze_filter else 0}_pn_{1 if change_nonlin else 0}/'
    Path(weight_dir).mkdir(exist_ok=True, parents=True)
    Path(stats_dir).mkdir(exist_ok=True, parents=True)
    cell_file = config_dict["sta_file"]
    cell_file = create_cell_file_from_config_version(
        original_file=cell_file, cell_index=cell_index, retina_index=retina
    )
    sta_dir = config_dict["stas_path"]
    receptive_field, max_coordinate = get_cell_sta(
        cell_file=cell_file,
        data_dir=f'{data_dir}/{sta_dir}',
        h=config['img_h'],
        w=config['img_w'],
        rf_size=(config['model_config']['num_of_channels'], config['rf_size'], config['rf_size']),
        flip_sta=flip_sta
    )

    for name, parameter in model.named_parameters():
        print(name)
        if not ('nonlinearity' in name):
            if not freeze_filter:
                if 'conv2' in name:
                    parameter.requires_grad = True
                    print('not freezing', name)
                else:
                    parameter.requires_grad = False
                    print('freezing', name)
        elif 'nonlinearity' in name:
            if not change_nonlin:
                parameter.requires_grad = False
                print('freezing', name)
            else:
                print('not freezing', name)
    print(model)

    print(config['model_config'].keys())

    for name, parameter in model.named_parameters():
        print(name)
    print(config.keys())

    if lr is not None:
        config['optimizer_config']['lr'] = lr
    else:
        lr = config['optimizer_config']['lr']
    if optimizer is None:
        optimizer = optim.Adam(model.parameters(), **config['optimizer_config'])
    if scheduler is None:
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, **config['scheduler_config'])
    if loss_function is None:
        loss_function = config['loss_function']
        loss_function = loss_function(**config["loss_function_config"])

    config['weights_dir'] = weight_dir
    config['stats_dir'] = stats_dir
    config['model_dir'] = model_dir
    config['freeze_filter'] = int(freeze_filter)
    config['finetuning_nonlin'] = lr
    config['parametrize_nonlin'] = change_nonlin

    if wandb_log:
        wandb.init(
            config=config,
            project=f"ln-factorized-{data_type}-{setting}",
            dir=log_dir,
            entity="retinal-circuit-modeling",
        )

    train(epochs=1000,
          dataloaders=dataloaders,
          dataset_config=config['dataloader_config'],
          config_dict=config_dict,
          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=stopper_patience,
          save_every_epoch=save_every_epoch,
          size=config['size'],
          max_coordinate=max_coordinate,
          wandb_log=wandb_log,
          fancy_nonlin=True,
          learn_filter=learn_filter,
          performance='validation'
    )
