import torch
import torch.nn as nn
from torch.nn.init import xavier_uniform_
from EMRI_DET.utilities import get_script_path
import dill as pickle
from pathlib import Path


class LinearModel(nn.Module):
    def __init__(self, in_features, out_features, neurons, n_layers, activation, name, initialisation=xavier_uniform_):
        super().__init__()
        self.initial = initialisation
        self.name = name

        layers = [nn.Linear(in_features, neurons[0]), activation()]
        for i in range(n_layers - 1):
            layers.append(nn.Linear(neurons[i], neurons[i + 1]))
            layers.append(activation())
        layers.append(nn.Linear(neurons[-1], out_features))

        self.layers = nn.Sequential(*layers)
        self.layers.apply(self.init_weights)

    def forward(self, x):
        return self.layers(x)

    def init_weights(self, m):
        if isinstance(m, nn.Linear):
            self.initial(m.weight)


def create_mlp(input_features, output_features, neurons, layers, activation, model_name, init=xavier_uniform_, device=None):
    if layers < 2:
        raise RuntimeError('You need at least two layers for an MLP.')
    if isinstance(neurons, list):
        if len(neurons) != layers:
            raise RuntimeError('Length of neuron vector does not equal number of hidden layers.')
    else:
        neurons = [neurons, ]
    model = LinearModel(input_features, output_features, neurons, layers, activation, model_name, initialisation=init)

    Path(get_script_path()+f'/../models/{model_name}/').mkdir(parents=True, exist_ok=True)
    pickle.dump(model, open(get_script_path()+f'/../models/{model_name}/function.pickle', "wb"), pickle.HIGHEST_PROTOCOL)  # save blank model

    if device is not None:
        model = model.to(device)
    return model


def load_mlp(model_name, get_state_dict=False):
    model = pickle.load(open(get_script_path()+f'/../models/{model_name}/function.pickle', "rb"))  # load blank model
    if get_state_dict:
        model.load_state_dict(torch.load(open(get_script_path()+f'/../models/{model_name}/model.pth')))
    return model
