import os

from helpers.data_generator import Sequence

from abc import ABC, abstractmethod, abstractproperty
import tensorflow as tf


gpus = tf.config.list_physical_devices("GPU")

if len(gpus):
    tf.config.set_logical_device_configuration(
        gpus[0], [tf.config.LogicalDeviceConfiguration(memory_limit=8072)]
    )

    logical_gpus = tf.config.list_logical_devices("GPU")


class BaseNetwork(ABC):
    @abstractmethod
    def get_network(self):
        pass

    @abstractproperty
    def input_name(self):
        pass


class BaseCombiner(ABC):
    @abstractmethod
    def get_output(self, input):
        pass


class NetworkConvNetDP(BaseNetwork):
    def __init__(self, series_shape, input_name="X_img", d_prob=0.2):
        self._series_shape = series_shape
        self._d_prob = d_prob
        self._input_name = input_name

    @property
    def input_name(self):
        return self._input_name

    def get_network(self, hp=None):
        input = tf.keras.layers.Input(self._series_shape, name=self.input_name)

        x = tf.keras.layers.Conv1D(
            2, 32, padding="same", activation="relu"
        )(input)
        x = tf.keras.layers.SpatialDropout1D(0)(x)
        x = tf.keras.layers.MaxPooling1D(2)(x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Conv1D(
            4, 10, padding="same", activation="relu"
        )(x)
        x = tf.keras.layers.SpatialDropout1D(0)(x)
        x = tf.keras.layers.MaxPooling1D(2)(x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Conv1D(
            4, 10, padding="same", activation="relu"
        )(x)
        x = tf.keras.layers.SpatialDropout1D(0)(x)
        x = tf.keras.layers.MaxPooling1D(2)(x)
        x = tf.keras.layers.BatchNormalization()(x)
        output = tf.keras.layers.Flatten()(x)

        return tf.keras.models.Model(input, output, name="img_net")


class NetworkConvNetFS(BaseNetwork):
    def __init__(self, img_shape, input_name="X_img"):
        self._img_shape = img_shape
        self._input_name = input_name

    @property
    def input_name(self):
        return self._input_name

    def get_network(self, hp=None):
        input = tf.keras.layers.Input(self._img_shape, name=self.input_name)
        flip_output = tf.keras.layers.RandomFlip(mode="horizontal")(input)

        x = tf.keras.layers.Conv2D(
            32,
            kernel_size=(2, 2),
            padding="same",
            activation="relu",
        )(flip_output)
        x = tf.keras.layers.SpatialDropout2D(0.517)(x)
        x = tf.keras.layers.MaxPooling2D((2, 2))(x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Conv2D(
            32, (3, 3), padding="same", activation="relu"
        )(x)
        x = tf.keras.layers.SpatialDropout2D(0.517)(x)
        x = tf.keras.layers.MaxPooling2D((2, 2))(x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Conv2D(
            10, (3, 3), padding="same", activation="relu"
        )(x)
        x = tf.keras.layers.SpatialDropout2D(0.517)(x)
        x = tf.keras.layers.MaxPooling2D((2, 2))(x)
        x = tf.keras.layers.BatchNormalization()(x)

        output = tf.keras.layers.Flatten()(x)

        return tf.keras.models.Model(input, output, name="img_net")


class VanillaMetNet(BaseNetwork):
    def __init__(self, input_shape, input_name="X_met", d_prob=0.2):
        self._input_shape = input_shape
        self._d_prob = d_prob
        self._input_name = input_name

    @property
    def input_name(self):
        return self._input_name

    def get_network(self, **kwargs):
        input = tf.keras.layers.Input(self._input_shape, name=self.input_name)
        output = input

        return tf.keras.models.Model(input, output, name="met_net")


class CombinerConvNetDP(BaseCombiner):
    def __init__(self):
        pass

    def get_output(self, input, hp=None):
        x = tf.keras.layers.Dense(512, activation="relu")(input)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Dropout(0.75)(x)
        x = tf.keras.layers.Dense(512, activation="relu")(x)
        x = tf.keras.layers.BatchNormalization()(x)
        output = tf.keras.layers.Dropout(0.75)(x)

        return output


class CombinerRefNetMD(BaseCombiner):
    def __init__(self):
        pass

    def get_output(self, input, hp=None):
        x = tf.keras.layers.Dense(512, activation="relu")(input)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Dropout(0.52485)(x)
        x = tf.keras.layers.Dense(512, activation="relu")(x)
        x = tf.keras.layers.BatchNormalization()(x)
        output = tf.keras.layers.Dropout(0.52485)(x)

        return output


class CombinerConvNetFS(BaseCombiner):
    def __init__(self):
        pass

    def get_output(self, input, hp=None):

        x = tf.keras.layers.Dense(256, activation="relu")(input)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Dropout(0.38)(x)
        x = tf.keras.layers.Dense(256, activation="relu")(x)
        x = tf.keras.layers.BatchNormalization()(x)
        output = tf.keras.layers.Dropout(0.38)(x)

        return output


class BaseOutputLayer(ABC):
    @abstractmethod
    def get_output(self, input):
        pass

    @abstractmethod
    def get_losses(self):
        pass

    @abstractmethod
    def get_metrics(self):
        pass


class VanillaOutput(BaseOutputLayer):
    def get_output(self, input):
        output = tf.keras.layers.Dense(1, activation="linear", name="vanilla_output")(
            input
        )
        return output

    def get_losses(self):
        return {"vanilla_output": "mse"}

    def get_metrics(self):
        return {
            "vanilla_output": [
                "mae",
                tf.keras.metrics.RootMeanSquaredError(name="RMSE"),
            ]
        }


class CompleteModel:
    def __init__(
        self,
        name: str,
        img_net: BaseNetwork,
        met_net: BaseNetwork,
        comb_net: BaseCombiner,
        out_layer: BaseOutputLayer,
        optimizer="Adam",
        base_path="./model_checkpoints/",
    ):
        self._name = name

        self._img_net = img_net
        self._met_net = met_net
        self._comb_net = comb_net
        self._out_layer = out_layer

        self._optimizer = optimizer
        self.base_path = base_path


    def build(self, hp=None):
        if self._img_net:
            img_net_inst = self._img_net.get_network(hp=hp)
        if self._met_net:
            met_net_inst = self._met_net.get_network(hp=hp)

        if self._img_net and self._met_net:
            comb_out = self._comb_net.get_output(
                tf.keras.layers.Concatenate()(
                    [img_net_inst.output, met_net_inst.output]
                ),
                hp=hp,
            )
        elif self._img_net:
            comb_out = self._comb_net.get_output(
                tf.keras.layers.Concatenate()([img_net_inst.output]), hp=hp
            )
        else:
            comb_out = self._comb_net.get_output(
                tf.keras.layers.Concatenate()([met_net_inst.output]), hp=hp
            )

        outputs = self._out_layer.get_output(comb_out)

        if self._img_net and self._met_net:
            self._model = tf.keras.models.Model(
                inputs={
                    img_net_inst.input.name: img_net_inst.input,
                    met_net_inst.input.name: met_net_inst.input,
                },
                outputs=outputs,
            )
        elif self._img_net:
            self._model = tf.keras.models.Model(
                inputs={img_net_inst.input.name: img_net_inst.input}, outputs=outputs
            )
        else:
            self._model = tf.keras.models.Model(
                inputs={met_net_inst.input.name: met_net_inst.input}, outputs=outputs
            )

        self._model.compile(
            loss=self._out_layer.get_losses(),
            metrics=self._out_layer.get_metrics(),
            optimizer=self._optimizer(hp) if hp else self._optimizer,
        )

        if hp:
            return self._model

        return self

    def load(self, checkpoint_path: str):
        self._model.load_weights(checkpoint_path)

        return self

    def save(self, checkpoint_path: str):
        self._model.save_weights(checkpoint_path)

        return self