from dataclasses import dataclass
from pathlib import Path
from typing import Union, List, Tuple

import torch
import torch.nn.functional as F
from huggingface_hub import ModelHubMixin
from torch import nn


class Dense(nn.Module):

    def __init__(self, input_dim, output_dim, bias=True, activation=nn.LeakyReLU, **kwargs):
        super().__init__()
        self.fc = nn.Linear(input_dim, output_dim, bias=bias)
        nn.init.xavier_uniform_(self.fc.weight)
        nn.init.constant_(self.fc.bias, 0.0)
        self.activation = activation(**kwargs) if activation is not None else None

    def forward(self, x):
        if self.activation is None:
            return self.fc(x)
        return self.activation(self.fc(x))


class Encoder(nn.Module):
    def __init__(self, input_dim, *dims):
        super().__init__()
        dims = (input_dim,) + dims
        self.layers = nn.Sequential(
            *[Dense(dims[i], dims[i+1], negative_slope=0.4, inplace=True) for i in range(len(dims) - 1)]
        )
    def forward(self, x):
        return self.layers(x)


class Decoder(nn.Module):
    def __init__(self, output_dim, *dims):
        super().__init__()
        self.layers = nn.Sequential(
            *[Dense(dims[i], dims[i + 1], negative_slope=0.4, inplace=True) for i in range(len(dims) - 1)]
            + [Dense(dims[-1], output_dim, activation=nn.Sigmoid)]
        )
    def forward(self, x):
        return self.layers(x)


@dataclass
class AutoencoderConfig:
    input_dim: int = 784
    hidden_dims: Union[Tuple[str], List[str]] = (256, 64, 16, 4, 2)


class Autoencoder(nn.Module, ModelHubMixin):

    def __init__(self, config: Union[dict, AutoencoderConfig] = AutoencoderConfig(), **kwargs):
        super().__init__()
        self.config = AutoencoderConfig(**config) if isinstance(config, dict) else config
        self.config.__dict__.update(**kwargs)

        self.encoder = Encoder(self.config.input_dim, *self.config.hidden_dims)
        self.decoder = Decoder(self.config.input_dim, *reversed(self.config.hidden_dims))

    def forward(self, x):
        x = x.flatten(1)
        latent = self.encoder(x)
        recon = self.decoder(latent)
        loss = F.mse_loss(recon, x)
        return recon, latent, loss

    def save_pretrained(self, save_directory, **kwargs):
        assert 'config' not in kwargs, \
            "save_pretrained handles passing model config for you, please dont pass it"
        super().save_pretrained(save_directory, config=self.config.__dict__, **kwargs)


class MessageModel:

    def __init__(self, msg='hello, world'):
        self.msg = msg

    def __call__(self):
        print(self.msg)

    @classmethod
    def from_pretrained(cls, path):
        path = Path(path)
        msg_file_path = path / 'message.txt'
        assert msg_file_path.exists()
        msg = msg_file_path.read_text()
        return cls(msg)

    def save_pretrained(self, path):
        path = Path(path)
        path.mkdir(exist_ok=True, parents=True)
        msg_file_path = path / 'message.txt'
        msg_file_path.write_text(self.msg)