import numpy as np
import torch

from . import blind_bwe_utils
from .sampler import BlindSampler


def apply_LTAS_init(recording, LTAS_ref, sample_rate, nfft):
    """
    Spectrally reshapes `recording` to match the checkpoint's reference
    long-term-average spectrum (LTAS), used both as the diffusion warm-start
    and as the reconstruction-guidance observation.

    From Evaluator.apply_LTAS_init, testing/evaluator.py:48-80.
    """
    freqs = np.fft.fftfreq(nfft, 1 / sample_rate)[0:nfft // 2 + 1]
    freqs[-1] = sample_rate / 2
    freqs = torch.tensor(freqs, dtype=torch.float32).to(recording.device)

    LTAS_smooth = blind_bwe_utils.smooth_LTAS(LTAS_ref, freqs, Noct=3)
    correction = 1 / LTAS_smooth.max()
    LTAS_smooth = (LTAS_smooth * correction).to(recording.device)

    win_length = nfft // 2
    hop_length = nfft // 4

    LTAS_y, std = blind_bwe_utils.compute_LTAS(recording.clone(), sample_rate, nfft=nfft, win_length=win_length, hop_length=hop_length, normalize=1)
    LTAS_y_smooth = blind_bwe_utils.smooth_LTAS(LTAS_y, freqs, Noct=3) * correction.to(LTAS_y.device)

    diff = torch.clamp(10 * torch.log10(LTAS_y_smooth) - 10 * torch.log10(LTAS_smooth), min=-20)

    X = torch.stft(recording, n_fft=nfft, hop_length=hop_length, win_length=win_length, window=torch.hann_window(win_length), return_complex=True)
    diff_lin = torch.pow(10, diff / 20)
    X = X / diff_lin.unsqueeze(-1)

    x_init = torch.istft(X, n_fft=nfft, hop_length=hop_length, win_length=win_length, window=torch.hann_window(win_length), length=recording.shape[0])
    x_init = std * x_init / x_init.std()
    return x_init


def segment_signal(recording, seg_len, overlap, device):
    """From Evaluator.segment_signal, testing/evaluator.py:82-99."""
    L = recording.shape[-1]
    ix_start = []
    ix_end = []
    segs = []
    ix = 0
    while ix + seg_len < L:
        segs.append(torch.Tensor(recording[..., ix:ix + seg_len]))
        ix_start.append(ix)
        ix_end.append(ix + seg_len)
        ix += seg_len - overlap

    segs.append(torch.cat((torch.Tensor(recording[..., ix::]).to(device), torch.zeros((seg_len - (L - ix),), device=device)), -1))
    ix_start.append(ix)
    ix_end.append(ix + seg_len)
    expanded_size = recording.shape[-1] + seg_len - (L - ix)
    return segs, ix_start, ix_end, expanded_size


def _crossfade_into_result(pred, y_masked, mask, seg_len, device):
    """
    Factored from the inline Hann-crossfade block inside the original's
    "Block-Autoregressive" branch, testing/evaluator.py:256-274.
    """
    if mask[0] == 1:
        first_0 = torch.where(mask == 0)[0][0]
        hann_window = torch.hann_window(int(first_0) * 2, device=device)
        hann_left = hann_window[0:first_0]
        hann_right = hann_window[first_0:]
        pred[..., 0:first_0] = pred[..., 0:first_0] * hann_left + y_masked[0:first_0] * hann_right

    if mask[-1] == 1:
        last_0 = torch.where(mask == 0)[0][-1]
        size = seg_len - last_0
        hann_window = torch.hann_window(int(size) * 2, device=device)
        hann_left = hann_window[0:size]
        hann_right = hann_window[size:]
        pred[..., seg_len - size:] = pred[..., seg_len - size:] * hann_right + y_masked[..., seg_len - size:] * hann_left

    return pred


def restore_audio(recording, cfg, network, diff_params, LTAS_ref, device, on_progress=None):
    """
    Restores a mono waveform (1-D tensor, already resampled to cfg.exp.sample_rate)
    by segmenting it into overlapping chunks and processing them autoregressively,
    each conditioned on the previous chunk's already-reconstructed overlap
    (Hann-crossfaded at the boundary). LTAS_init/LTAS_as_y are always True here,
    so the LTAS-corrected signal doubles as both the diffusion warm-start and the
    reconstruction-guidance observation for every segment.

    From the "Block-Autoregressive" branch of Evaluator.evaluate_single_recording,
    testing/evaluator.py:101-347 (lines 229-283).
    """
    sample_rate = cfg.exp.sample_rate
    seg_len = cfg.tester.evaluation.segment_length
    overlap = int(cfg.tester.evaluation.overlap * seg_len)
    sigma_norm = cfg.tester.blind_bwe.sigma_norm

    recording = apply_LTAS_init(recording, LTAS_ref, sample_rate, cfg.tester.blind_bwe.LTAS_fft)

    segs, ix_start, ix_end, expanded_size = segment_signal(recording, seg_len, overlap, device)

    std_orig = torch.median(torch.Tensor([seg.std() for seg in segs]))
    segs = [(sigma_norm * seg / std_orig).to(torch.float32) for seg in segs]

    sampler = BlindSampler(model=network, diff_params=diff_params, args=cfg)

    n_segments = len(segs)
    result = torch.zeros((expanded_size,), device="cpu")
    result_mask = torch.zeros((expanded_size,), device="cpu")

    def step_cb(seg_idx):
        def cb(step_i, total_steps):
            if on_progress is not None:
                frac = (seg_idx + (step_i + 1) / total_steps) / n_segments
                on_progress(frac, f"segment {seg_idx + 1}/{n_segments}, step {step_i + 1}/{total_steps}")
        return cb

    y0 = segs[0].unsqueeze(0).to(device)
    pred, _ = sampler.predict_blind_bwe(y0, x_init=y0, progress_cb=step_cb(0))
    result[ix_start[0]:ix_end[0]] = pred[0].cpu()
    result_mask[ix_start[0]:ix_end[0]] = 1

    for i in range(1, n_segments):
        y = segs[i].to(device)
        x_init = segs[i].to(device)
        y_masked = result[ix_start[i]:ix_end[i]].clone().to(device)
        mask = result_mask[ix_start[i]:ix_end[i]].clone().to(device)

        pred, _ = sampler.predict_blind_bwe_AR(y, y_masked, mask=mask, x_init=x_init, progress_cb=step_cb(i))
        pred = _crossfade_into_result(pred[0], y_masked, mask, seg_len, device)

        result_mask[ix_start[i]:ix_end[i]] = 1
        result[ix_start[i]:ix_end[i]] = pred.cpu()

    result = result[0:recording.shape[-1]]
    result = std_orig * result / sigma_norm
    return result