import logging
import os
import shutil

import numpy as np
import pandas as pd
import torch
from glm_saga.elasticnet import glm_saga
from torch import nn

from sparsification.FeatureSelection import FeatureSelectionFitting
from sparsification import data_helpers
from sparsification.utils import get_default_args, compute_features_and_metadata, select_in_loader, get_feature_loaders


def get_glm_selection(feature_loaders, metadata,  args, num_classes, device, n_features_to_select, folder):
    num_features = metadata["X"]["num_features"][0]
    fittingClass = FeatureSelectionFitting(num_features, num_classes, args, 0.8,
                                           n_features_to_select,
                                           0.1,folder,
                                           lookback=3, tol=1e-4,
                                           epsilon=1,)
    to_drop, test_acc = fittingClass.fit(feature_loaders, metadata, device)
    selected_features = torch.tensor([i for i in range(num_features) if i not in to_drop])
    return selected_features


def compute_feature_selection_and_assignment(model, train_loader, test_loader, log_folder,num_classes, seed,  select_features = 50):
    feature_loaders, metadata, device,args =  get_feature_loaders(seed, log_folder,train_loader, test_loader, model, num_classes, )

    if os.path.exists(log_folder / f"SlDD_Selection_{select_features}.pt"):
        feature_selection = torch.load(log_folder / f"SlDD_Selection_{select_features}.pt")
    else:
        used_features = model.linear.weight.shape[1]
        if used_features != select_features:
            selection_folder = log_folder / "sldd_selection" # overwrite with None to prevent saving
            feature_selection = get_glm_selection(feature_loaders, metadata,   args,
                                      num_classes,
                                      device,select_features, selection_folder
                                      )
        else:
            feature_selection = model.linear.selection
        torch.save(feature_selection, log_folder / f"SlDD_Selection_{select_features}.pt")
    feature_loaders = select_in_loader(feature_loaders, feature_selection)
    mean, std = metadata["X"]["mean"], metadata["X"]["std"]
    mean_to_pass_in = mean
    std_to_pass_in = std
    if len(mean) != feature_selection.shape[0]:
        mean_to_pass_in = mean[feature_selection]
        std_to_pass_in = std[feature_selection]

    sparse_matrices, biases = fit_glm(log_folder, mean_to_pass_in, std_to_pass_in, feature_loaders, num_classes, select_features)

    return feature_selection, sparse_matrices, biases, mean, std


def fit_glm(log_dir,mean, std , feature_loaders,  num_classes,   select_features = 50):
    output_folder = log_dir / "glm_path"
    if not output_folder.exists() or len(list(output_folder.iterdir())) != 102:
        shutil.rmtree(output_folder, ignore_errors=True)
        output_folder.mkdir(exist_ok=True, parents=True)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        linear = nn.Linear(select_features, num_classes).to(device)
        for p in [linear.weight, linear.bias]:
            p.data.zero_()
        print("Preparing normalization preprocess and indexed dataloader")
        metadata = {"X": {"mean": mean, "std": std},}
        preprocess = data_helpers.NormalizedRepresentation(feature_loaders['train'],
                                                           metadata=metadata,
                                                           device=linear.weight.device)

        print("Calculating the regularization path")
        mpl_logger = logging.getLogger("matplotlib")
        mpl_logger.setLevel(logging.WARNING)
        params = glm_saga(linear,
                          feature_loaders['train'],
                          0.1,
                          2000,
                          0.99, k=100,
                          val_loader=feature_loaders['val'],
                          test_loader=feature_loaders['test'],
                          n_classes=num_classes,
                          checkpoint=str(output_folder),
                          verbose=200,
                          tol=1e-4, # Change for ImageNet
                          lookbehind=5,
                          lr_decay_factor=1,
                          group=False,
                          epsilon=0.001,
                          metadata=None, # To let it be recomputed
                          preprocess=preprocess, )
    results = load_glm(output_folder)
    sparse_matrices = results["weights"]
    biases =  results["biases"]

    return sparse_matrices, biases

def load_glm(result_dir):
    Nlambda = max([int(f.split('params')[1].split('.pth')[0])
                   for f in os.listdir(result_dir) if 'params' in f]) + 1

    print(f"Loading regularization path of length {Nlambda}")

    params_dict = {i: torch.load(os.path.join(result_dir, f"params{i}.pth"),
                              map_location=torch.device('cpu')) for i in range(Nlambda)}

    regularization_strengths = [params_dict[i]['lam'].item() for i in range(Nlambda)]
    weights = [params_dict[i]['weight'] for i in range(Nlambda)]
    biases = [params_dict[i]['bias'] for i in range(Nlambda)]

    metrics = {'acc_tr': [], 'acc_val': [], 'acc_test': []}

    for k in metrics.keys():
        for i in range(Nlambda):
            metrics[k].append(params_dict[i]['metrics'][k])
        metrics[k] = 100 * np.stack(metrics[k])
    metrics = pd.DataFrame(metrics)
    metrics = metrics.rename(columns={'acc_tr': 'acc_train'})

    # weights_stacked = ch.stack(weights)
    # sparsity = ch.sum(weights_stacked != 0, dim=2).numpy()
    sparsity = np.array([torch.sum(w != 0, dim=1).numpy() for w in weights])

    return {'metrics': metrics,
            'regularization_strengths': regularization_strengths,
            'weights': weights,
            'biases': biases,
            'sparsity': sparsity,
            'weight_dense': weights[-1],
            'bias_dense': biases[-1]}