import sys

sys.path.insert(0, "/home/mikeys/src/hoot")
from dataclasses import dataclass
import inspect
from typing import Optional

import torch
import torch.nn.functional as F

from hoot.modules.featurizer import Featurizer
from hoot.modules.encoder import Encoder


@dataclass
class HootConfig:
    n_layers: int = 12
    n_embd: int = 256
    n_classes: int = 2
    highfreq: Optional[float] = None


class Decoder(torch.nn.Module):
    def __init__(self, feat_in, num_classes):
        super().__init__()
        self._feat_in = feat_in
        self._num_classes = num_classes
        # self.conv = torch.nn.Conv1d(self._feat_in, self._num_classes, kernel_size=1, bias=True)
        self.relu = torch.nn.ReLU()
        self.drop = torch.nn.Dropout(0.2)
        self.head = torch.nn.Linear(self._feat_in, self._num_classes, bias=True)

    def forward(self, encoder_output):
        x = encoder_output
        x = torch.mean(x, dim=-1)
        x = self.drop(x)
        x = self.relu(x)
        return self.head(x)


class Model(torch.nn.Module):
    def __init__(self, config=None):
        super(Model, self).__init__()
        self.config = config if config is not None else HootConfig()
        config = self.config
        self.featurizer = Featurizer(highfreq=config.highfreq)
        self.encoder = Encoder(config.n_layers, config.n_embd)
        self.decoder = Decoder(config.n_embd, config.n_classes)
        # init all weights
        self.apply(self._init_weights)

    def forward(self, signals, signals_len, targets=None, targets_len=None):
        assert (targets is None) == (targets_len is None)
        processed_signal, processed_signal_length = self.featurizer(
            input_signal=signals, length=signals_len
        )
        encoded, encoded_len = self.encoder(
            audio_signal=processed_signal, length=processed_signal_length
        )
        logits = self.decoder(encoder_output=encoded)
        loss = None
        if targets is not None:
            # loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
            loss = F.cross_entropy(logits, targets, ignore_index=-1)
        return logits, loss

    def _init_weights(self, module):
        if isinstance(module, torch.nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, torch.nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def eval(self):
        for name, module in self.named_modules():
            module.train("batch" in name.split(".")[-1])

    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type, use_fused=True):
        # start with all of the candidate parameters
        param_dict = {pn: p for pn, p in self.named_parameters()}
        # filter out those that do not require grad
        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
        # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
        optim_groups = [
            {"params": decay_params, "weight_decay": weight_decay},
            {"params": nodecay_params, "weight_decay": 0.0},
        ]
        num_decay_params = sum(p.numel() for p in decay_params)
        num_nodecay_params = sum(p.numel() for p in nodecay_params)
        print(
            f"num decayed parameter tensors: {len(decay_params)},"
            f" with {num_decay_params:,} parameters"
        )
        print(
            f"num non-decayed parameter tensors: {len(nodecay_params)},"
            f" with {num_nodecay_params:,} parameters"
        )
        # Create AdamW optimizer and use the fused version if it is available
        fused_available = "fused" in inspect.signature(torch.optim.AdamW).parameters
        extra_args = (
            dict(fused=True) if fused_available and device_type == "cuda" and use_fused else dict()
        )
        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
        print(f"using fused AdamW: {use_fused}")
        return optimizer