import abc
import json

import numpy
import scipy.linalg
from keras import Sequential
from keras.layers import Dense
from keras.models import load_model

model_types = {}


class Model:
    @abc.abstractmethod
    def get_temperature(self, flat_history: numpy.array) -> float:
        pass

    @abc.abstractmethod
    def save(self, filename):
        pass

    @classmethod
    def load(cls, config, filename):
        return model_types[config['model_type']]._load(filename)

    @abc.abstractclassmethod
    def _load(cls, filename):
        pass

    @classmethod
    def generate(cls, config: dict):
        return model_types[config['model_type']].generate(config)


class LinearModel(Model):
    def __init__(self, model):
        self.model = model

    def get_temperature(self, flat_history):
        return sum(numpy.multiply(self.model, flat_history))

    def save(self, filename):
        with open(filename, 'w') as fp:
            json.dump(self.model.tolist(), fp)

    @classmethod
    def _load(cls, filename):
        with open(filename, 'r') as fp:
            model = json.load(fp)
        model = numpy.array(model)
        return cls(model)

    @classmethod
    def generate(cls, config: dict):
        from .ArgParser import get_sliding_window_from_config
        sliding_window = get_sliding_window_from_config(config)
        xs = []
        ys = []
        for x in sliding_window:
            xs.append(x.get_model_values())
            ys.append(x.get_model_target())
        return cls(scipy.linalg.lstsq(xs, ys)[0])


model_types['linear'] = LinearModel


class NeuralModel(Model):
    def __init__(self, model):
        self.model = model

    def get_temperature(self, flat_history):
        state = numpy.reshape(flat_history, [1, len(flat_history)])
        result = self.model.predict(state)
        return result[0][0]

    @staticmethod
    def _build_model(state_size):
        model = Sequential()
        model.add(Dense(int(state_size / 2) or 1, activation='linear',
                        input_dim=state_size,
                        kernel_initializer='random_uniform'))
        model.add(Dense(int(state_size / 4) or 1, activation='linear',
                        kernel_initializer='random_uniform'))
        model.add(Dense(int(state_size / 8) or 1, activation='linear',
                        kernel_initializer='random_uniform'))
        model.add(Dense(1, activation='linear',
                        kernel_initializer='random_uniform'))
        model.compile(optimizer='adagrad',
                      loss='mean_squared_error',
                      metrics=['accuracy'])
        return model

    def save(self, filename):
        self.model.save(filename)

    @classmethod
    def _load(cls, filename):
        model = load_model(filename)
        return cls(model)

    @classmethod
    def generate(cls, config: dict):
        from .ArgParser import get_sliding_window_from_config
        sliding_window = get_sliding_window_from_config(config)
        window = next(sliding_window)
        model = cls._build_model(window.get_model_size())

        xs = []
        ys = []
        for window in sliding_window:
            xs.append(window.get_model_values())
            ys.append([window.get_next_value('temp_in')])

        model.fit(numpy.array(xs), numpy.array(ys), epochs=config['model_epochs'], batch_size=32, shuffle=True)
        return cls(model)


model_types['neural'] = NeuralModel


class NeuralModelLess(NeuralModel):
    @staticmethod
    def _build_model(state_size):
        model = Sequential()
        model.add(Dense(int(state_size / 2) or 1, activation='linear',
                        input_dim=state_size,
                        kernel_initializer='random_uniform'))
        model.add(Dense(int(state_size / 4) or 1, activation='linear',
                        kernel_initializer='random_uniform'))
        model.add(Dense(int(state_size / 8) or 1, activation='linear',
                        kernel_initializer='random_uniform'))
        model.add(Dense(1, activation='linear',
                        kernel_initializer='random_uniform'))
        model.compile(optimizer='adagrad',
                      loss='mean_squared_error',
                      metrics=['accuracy'])
        return model

model_types['neural_less'] = NeuralModelLess


class NeuralModelMore(NeuralModel):
    @staticmethod
    def _build_model(state_size):
        model = Sequential()
        model.add(Dense(int(state_size / 2) or 1, activation='linear',
                        input_dim=state_size,
                        kernel_initializer='random_uniform'))
        model.add(Dense(int(state_size / 4) or 1, activation='linear',
                        kernel_initializer='random_uniform'))
        model.add(Dense(int(state_size / 8) or 1, activation='linear',
                        kernel_initializer='random_uniform'))
        model.add(Dense(1, activation='linear',
                        kernel_initializer='random_uniform'))
        model.compile(optimizer='adagrad',
                      loss='mean_squared_error',
                      metrics=['accuracy'])
        return model

model_types['neural_more'] = NeuralModelMore


class NeuralModelRelu(NeuralModel):
    @staticmethod
    def _build_model(state_size):
        model = Sequential()
        model.add(Dense(int(state_size / 2) or 1, activation='linear',
                        input_dim=state_size,
                        kernel_initializer='random_uniform'))
        model.add(Dense(int(state_size / 4) or 1, activation='relu',
                        kernel_initializer='random_uniform'))
        model.add(Dense(int(state_size / 8) or 1, activation='relu',
                        kernel_initializer='random_uniform'))
        model.add(Dense(1, activation='linear',
                        kernel_initializer='random_uniform'))
        model.compile(optimizer='adagrad',
                      loss='mean_squared_error',
                      metrics=['accuracy'])
        return model

model_types['neural_relu'] = NeuralModelRelu


class NeuralModelReduced(NeuralModel):
    @staticmethod
    def _build_model(state_size):
        model = Sequential()
        model.add(Dense(int(state_size / 2) or 1, activation='linear',
                        input_dim=state_size,
                        kernel_initializer='random_uniform'))
        model.add(Dense(1, activation='linear',
                        kernel_initializer='random_uniform'))
        model.compile(optimizer='adagrad',
                    loss='mean_squared_error',
                    metrics=['accuracy'])
        return model

model_types['neural_half'] = NeuralModelReduced


class NeuralModelSqueezeExpand(NeuralModel):
    @staticmethod
    def _build_model(state_size):
        model = Sequential()
        model.add(Dense(int(state_size / 2) or 1, activation='linear',
                        input_dim=state_size,
                        kernel_initializer='random_uniform'))
        model.add(Dense(int(state_size / 4) or 1, activation='linear',
                        kernel_initializer='random_uniform'))
        model.add(Dense(int(state_size / 2) or 1, activation='linear',
                        kernel_initializer='random_uniform'))
        model.add(Dense(1, activation='linear',
                        kernel_initializer='random_uniform'))
        model.compile(optimizer='adagrad',
                      loss='mean_squared_error',
                      metrics=['accuracy'])
        return model

model_types['neural_hqh'] = NeuralModelSqueezeExpand