from __future__ import annotations

import os
import threading
import urllib.request
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torchaudio

SAMPLE_RATE = 44100
CACHE_ROOT = Path(
    os.environ.get("MICROTCN_CACHE_DIR", Path.home() / ".cache" / "micro-tcn")
)

CKPT_URL = (
    "https://raw.githubusercontent.com/csteinmetz1/neural-2a/master/models/"
    "traced_3-uTCN-300__causal__4-10-13__fraction-1.0-bs32.pt"
)
CKPT_FILENAME = "traced_3-uTCN-300__causal__4-10-13__fraction-1.0-bs32.pt"

_LOCK = threading.Lock()
_MODEL = None


class FiLM(nn.Module):
    def __init__(self, num_features: int, cond_dim: int):
        super().__init__()
        self.bn = nn.BatchNorm1d(num_features, affine=False)
        self.adaptor = nn.Linear(cond_dim, num_features * 2)

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        cond = self.adaptor(cond)
        g, b = torch.chunk(cond, 2, dim=-1)
        g = g.permute(0, 2, 1)
        b = b.permute(0, 2, 1)
        return g * self.bn(x) + b


class TCNBlock(nn.Module):
    def __init__(
        self,
        in_ch: int,
        out_ch: int,
        kernel_size: int,
        dilation: int,
        cond_dim: int,
        causal: bool = True,
    ):
        super().__init__()
        self.causal = causal
        self.conv1 = nn.Conv1d(
            in_ch, out_ch,
            kernel_size=kernel_size,
            dilation=dilation,
            padding=0,
            bias=False,
        )
        self.film = FiLM(out_ch, cond_dim)
        self.relu = nn.PReLU(out_ch)
        self.res = nn.Conv1d(in_ch, out_ch, kernel_size=1, groups=in_ch, bias=False)

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        h = self.conv1(x)
        h = self.film(h, cond)
        h = self.relu(h)
        res = self.res(x)
        # Match the crop used by the author's released TCN implementation.
        if self.causal:
            stop = res.shape[-1] - 1
            res = res[..., stop - h.shape[-1]:stop]
        else:
            start = (res.shape[-1] - h.shape[-1]) // 2
            res = res[..., start:start + h.shape[-1]]
        return h + res


class MicroTCN(nn.Module):
    def __init__(
        self,
        nparams: int = 2,
        nblocks: int = 4,
        kernel_size: int = 13,
        dilation_growth: int = 10,
        channel_width: int = 32,
        causal: bool = True,
    ):
        super().__init__()
        self.nparams = nparams
        cond_dim = channel_width

        self.gen = nn.Sequential(
            nn.Linear(nparams, 16),
            nn.ReLU(),
            nn.Linear(16, cond_dim),
            nn.ReLU(),
            nn.Linear(cond_dim, cond_dim),
            nn.ReLU(),
        )

        self.blocks = nn.ModuleList()
        for n in range(nblocks):
            in_ch = 1 if n == 0 else channel_width
            dilation = dilation_growth ** n
            self.blocks.append(
                TCNBlock(in_ch, channel_width, kernel_size, dilation, cond_dim, causal)
            )

        self.output = nn.Conv1d(channel_width, 1, kernel_size=1)

    def forward(self, x: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
        cond = self.gen(p)
        for block in self.blocks:
            x = block(x, cond)
        return torch.tanh(self.output(x))


def _download_checkpoint() -> Path:
    if checkpoint := os.environ.get("MICROTCN_CHECKPOINT"):
        return Path(checkpoint).expanduser()

    CACHE_ROOT.mkdir(parents=True, exist_ok=True)
    ckpt_path = CACHE_ROOT / CKPT_FILENAME

    if ckpt_path.exists():
        return ckpt_path

    partial = Path(f"{ckpt_path}.part")
    partial.unlink(missing_ok=True)
    print("Downloading micro-TCN checkpoint...", flush=True)
    request = urllib.request.Request(
        CKPT_URL,
        headers={"User-Agent": "TEAMuP-micro-TCN/1.0"},
    )
    try:
        with (
            urllib.request.urlopen(request, timeout=120) as response,
            partial.open("wb") as f,
        ):
            while chunk := response.read(1 << 16):
                f.write(chunk)
        partial.replace(ckpt_path)
    except Exception:
        partial.unlink(missing_ok=True)
        raise

    return ckpt_path


def _load_model() -> MicroTCN:
    global _MODEL
    if _MODEL is not None:
        return _MODEL

    with _LOCK:
        if _MODEL is not None:
            return _MODEL

        ckpt_path = _download_checkpoint()
        traced = torch.jit.load(ckpt_path, map_location="cpu")
        state = traced.state_dict()

        model = MicroTCN(
            nparams=2,
            nblocks=4,
            kernel_size=13,
            dilation_growth=10,
            channel_width=32,
            causal=True,
        )
        model.load_state_dict(state, strict=True)
        model.eval()
        _MODEL = model

    return _MODEL


_RECEPTIVE_FIELD = None


def _receptive_field(model, device) -> int:
    """
    How many samples the network consumes, measured once from the model itself.

    Args:
        model: The loaded network.
        device: Where to run the measurement.

    Returns:
        deficit (int): Samples the output is short of the input, 0 if none.
    """

    global _RECEPTIVE_FIELD

    if _RECEPTIVE_FIELD is None:
        probe = torch.zeros(1, 1, SAMPLE_RATE, device=device)
        params = torch.zeros(1, 1, 2, device=device)
        with torch.inference_mode():
            out = model(probe, params)
        _RECEPTIVE_FIELD = max(0, probe.shape[-1] - out.shape[-1])

    return _RECEPTIVE_FIELD


def process_audio(
    audio_path: str,
    limit_mode: float,
    peak_reduction: float,
) -> tuple[np.ndarray, int]:
    model = _load_model()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    waveform, sr = torchaudio.load(audio_path)
    waveform = waveform.mean(dim=0, keepdim=True).float()
    if sr != SAMPLE_RATE:
        waveform = torchaudio.functional.resample(waveform, sr, SAMPLE_RATE)

    # The network does not pad, so it returns its receptive field fewer samples
    # than it is given. Left-padding by that deficit, with the first sample
    # repeated so the pad is not a step into silence, gives a result the same
    # length as the input and aligned with it.
    pad = _receptive_field(model, device)
    x = waveform.unsqueeze(0).to(device)

    if pad > 0:
        x = torch.cat([x[..., :1].expand(-1, -1, pad), x], dim=-1)

    p = torch.tensor(
        [[[limit_mode, peak_reduction]]], dtype=torch.float32, device=device
    )

    with torch.inference_mode():
        y = model(x, p)

    output = y.squeeze().cpu().numpy().astype(np.float32)

    # Guard against a deficit that is not exactly the measured one
    wanted = waveform.shape[-1]
    if output.shape[-1] > wanted:
        output = output[..., :wanted]

    return output, SAMPLE_RATE