from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn
from torch.nn.utils import weight_norm

try:
    from .config import ExperimentConfig, ModelConfig
except ImportError:
    from config import ExperimentConfig, ModelConfig


class Chomp1d(nn.Module):
    def __init__(self, chomp_size: int) -> None:
        super().__init__()
        self.chomp_size = chomp_size

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.chomp_size == 0:
            return x
        return x[:, :, :-self.chomp_size].contiguous()


class CausalConv1d(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, dilation: int = 1) -> None:
        super().__init__()
        padding = (kernel_size - 1) * dilation
        self.net = nn.Sequential(
            weight_norm(
                nn.Conv1d(
                    in_channels=in_channels,
                    out_channels=out_channels,
                    kernel_size=kernel_size,
                    padding=padding,
                    dilation=dilation,
                )
            ),
            Chomp1d(padding),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class TemporalBlock(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, dilation: int, dropout: float) -> None:
        super().__init__()
        self.conv1 = CausalConv1d(in_channels, out_channels, kernel_size, dilation=dilation)
        self.act1 = nn.GELU()
        self.dropout1 = nn.Dropout(dropout)
        self.conv2 = CausalConv1d(out_channels, out_channels, kernel_size, dilation=dilation)
        self.act2 = nn.GELU()
        self.dropout2 = nn.Dropout(dropout)
        self.residual = nn.Conv1d(in_channels, out_channels, kernel_size=1) if in_channels != out_channels else nn.Identity()
        self.final_act = nn.GELU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = self.residual(x)
        out = self.dropout1(self.act1(self.conv1(x)))
        out = self.dropout2(self.act2(self.conv2(out)))
        return self.final_act(out + residual)


@dataclass
class ReceptiveFieldInfo:
    receptive_field: int
    dilations: tuple[int, ...]


class RMSRegressionTCN(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        blocks: list[nn.Module] = []
        in_channels = config.input_channels
        dilations: list[int] = []

        for level, out_channels in enumerate(config.tcn_channels):
            dilation = config.dilation_base ** level
            dilations.append(dilation)
            blocks.append(
                TemporalBlock(
                    in_channels=in_channels,
                    out_channels=out_channels,
                    kernel_size=config.kernel_size,
                    dilation=dilation,
                    dropout=config.dropout,
                )
            )
            in_channels = out_channels

        self.encoder = nn.Sequential(*blocks)
        self.rms_head = nn.Sequential(
            nn.Linear(in_channels * 2 + 3, config.pooled_feature_dim),
            nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.pooled_feature_dim, config.pooled_feature_dim),
            nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.pooled_feature_dim, 1),
        )
        self.waveform_head = nn.Sequential(
            nn.Conv1d(in_channels, in_channels, kernel_size=1),
            nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Conv1d(in_channels, 1, kernel_size=1),
        )
        self.receptive_field_info = compute_receptive_field(config)

    def forward(self, x: torch.Tensor, mask: torch.Tensor, aux: torch.Tensor) -> dict[str, torch.Tensor]:
        if x.ndim != 3:
            raise ValueError(f"Expected x shape (batch, seq, channels), got {tuple(x.shape)}")
        features = self.encoder(x.transpose(1, 2))
        mask_1d = mask.unsqueeze(1)
        masked_features = features * mask_1d
        valid_count = mask_1d.sum(dim=2).clamp_min(1.0)
        mean_pool = masked_features.sum(dim=2) / valid_count
        masked_for_max = features.masked_fill(mask_1d == 0.0, float("-inf"))
        max_pool = masked_for_max.max(dim=2).values
        max_pool = torch.where(torch.isfinite(max_pool), max_pool, torch.zeros_like(max_pool))
        fused = torch.cat([mean_pool, max_pool, aux], dim=1)
        rms_prediction = self.rms_head(fused)
        waveform_prediction = self.waveform_head(features).transpose(1, 2)
        return {"rms": rms_prediction, "waveform": waveform_prediction}


def compute_receptive_field(config: ModelConfig) -> ReceptiveFieldInfo:
    receptive_field = 1
    dilations: list[int] = []
    for level, _ in enumerate(config.tcn_channels):
        dilation = config.dilation_base ** level
        dilations.append(dilation)
        receptive_field += 2 * (config.kernel_size - 1) * dilation
    return ReceptiveFieldInfo(receptive_field=receptive_field, dilations=tuple(dilations))


def build_model(config: ExperimentConfig | ModelConfig) -> RMSRegressionTCN:
    model_config = config.model if isinstance(config, ExperimentConfig) else config
    return RMSRegressionTCN(model_config)