#!/usr/bin/env python3
"""
SDXL Long Context Converter — Learned Sinusoidal Extrapolation (Optimized)

Extends positional embeddings by fitting parametric sinusoidal models
to each dimension, matching the mathematical structure learned during
CLIP training.

Methods:
  learned_sinusoidal  — fit N sinusoids + trend
  vaswani_fitted      — fit Vaswani PE formula
  hybrid              — best of all per dimension (recommended)
"""

import argparse
import json
import os
import sys
import time

import numpy as np
import torch

try:
    from safetensors import safe_open
    from safetensors.torch import save_file
except ImportError:
    print("ERROR: pip install safetensors")
    sys.exit(1)

try:
    from scipy.optimize import curve_fit
except ImportError:
    print("Installing scipy...")
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "scipy"])
    from scipy.optimize import curve_fit


POS_EMB_KEYS = [
    "conditioner.embedders.0.transformer.text_model.embeddings.position_embedding.weight",
    "cond_stage_model.transformer.text_model.embeddings.position_embedding.weight",
    "text_encoder.text_model.embeddings.position_embedding.weight",
    "conditioner.embedders.1.model.positional_embedding",
    "conditioner.embedders.1.transformer.text_model.embeddings.position_embedding.weight",
    "text_encoder_2.text_model.embeddings.position_embedding.weight",
]

MAX_LENGTH = 248


# ================================================================
# Progress bar
# ================================================================

def print_progress(current, total, prefix="", width=40):
    pct = current / total
    filled = int(width * pct)
    bar = "█" * filled + "░" * (width - filled)
    print(f"\r  {prefix} [{bar}] {current}/{total} ({pct*100:.0f}%)", end="", flush=True)


# ================================================================
# Vectorized spectral fallback (fast)
# ================================================================

def spectral_extrapolate_single(signal, target_length, damping=0.08):
    """Fast spectral extrapolation for a single dimension."""
    orig_len = len(signal)
    dc = np.mean(signal)
    spectrum = np.fft.rfft(signal - dc)
    n_freqs = len(spectrum)

    t_new = np.arange(target_length, dtype=np.float64) / orig_len

    # Vectorized: compute all frequencies at once
    amps = np.abs(spectrum)
    phases = np.angle(spectrum)

    # Normalization factors
    norms = np.full(n_freqs, 2.0 / orig_len)
    norms[0] = 1.0 / orig_len
    if orig_len % 2 == 0:
        norms[-1] = 1.0 / orig_len

    # Damping envelope for extrapolated region
    extrap_dist = np.zeros(target_length, dtype=np.float64)
    extrap_dist[orig_len:] = np.arange(1, target_length - orig_len + 1, dtype=np.float64) / orig_len

    reconstructed = np.zeros(target_length, dtype=np.float64)

    # Vectorize over positions (not frequencies) for better cache behavior
    k_values = np.arange(n_freqs, dtype=np.float64)
    freq_ratios = k_values / n_freqs

    for k in range(n_freqs):
        if amps[k] < 1e-10:
            continue
        damp = np.exp(-damping * freq_ratios[k] * freq_ratios[k] * extrap_dist * orig_len)
        reconstructed += amps[k] * norms[k] * np.cos(2 * np.pi * k * t_new - phases[k]) * damp

    return reconstructed + dc


def spectral_extrapolate_batch(signals, target_length, damping=0.08):
    """Batch spectral extrapolation for multiple dimensions at once."""
    n_dims, orig_len = signals.shape
    result = np.zeros((n_dims, target_length), dtype=np.float64)

    for d in range(n_dims):
        result[d] = spectral_extrapolate_single(signals[d], target_length, damping)

    return result


# ================================================================
# Sinusoidal model fitting (optimized)
# ================================================================

def sinusoidal_model_2(pos, a, b, A1, f1, p1, A2, f2, p2):
    return (a + b * pos +
            A1 * np.sin(2 * np.pi * f1 * pos + p1) +
            A2 * np.sin(2 * np.pi * f2 * pos + p2))


def sinusoidal_model_3(pos, a, b, A1, f1, p1, A2, f2, p2, A3, f3, p3):
    return (a + b * pos +
            A1 * np.sin(2 * np.pi * f1 * pos + p1) +
            A2 * np.sin(2 * np.pi * f2 * pos + p2) +
            A3 * np.sin(2 * np.pi * f3 * pos + p3))


def fit_sinusoidal_fast(signal, positions, n_components=3):
    """
    Fast sinusoidal fitting with reduced maxfev and smart initialization.
    Returns (params, model_func, r_squared).
    """
    orig_len = len(signal)
    dc = np.mean(signal)
    trend = (signal[-1] - signal[0]) / orig_len if orig_len > 1 else 0.0

    spectrum = np.fft.rfft(signal - dc)
    amps = np.abs(spectrum)
    top_freqs = np.argsort(amps[1:])[::-1][:n_components] + 1

    if n_components <= 2:
        model_func = sinusoidal_model_2
        n_comp = 2
    else:
        model_func = sinusoidal_model_3
        n_comp = 3

    p0 = [dc, trend]
    for i in range(n_comp):
        idx = top_freqs[i] if i < len(top_freqs) else 1
        p0.extend([
            amps[idx] * 2 / orig_len,
            idx / orig_len,
            0.0
        ])

    try:
        popt, _ = curve_fit(model_func, positions, signal, p0=p0, maxfev=2000)
        fitted = model_func(positions, *popt)
        residual = np.mean((signal - fitted) ** 2)
        variance = np.var(signal)
        r2 = 1.0 - residual / (variance + 1e-10) if variance > 1e-10 else 1.0
        return popt, model_func, r2
    except (RuntimeError, ValueError):
        return None, model_func, 0.0


def fit_vaswani_fast(signal, dim_idx, d_model, positions):
    """
    Fast Vaswani formula fitting with reduced iterations.
    Returns (params, model_func, r_squared).
    """
    dc = np.mean(signal)
    signal_ac = signal - dc
    amp_init = np.std(signal_ac) * np.sqrt(2)

    def model(pos, amplitude, base, phase, offset):
        freq = 1.0 / np.power(np.abs(base) + 1.0, 2.0 * dim_idx / d_model)
        return offset + amplitude * np.sin(pos * freq + phase)

    try:
        popt, _ = curve_fit(
            model, positions, signal,
            p0=[amp_init, 10000.0, 0.0, dc],
            maxfev=2000,
            bounds=(
                [-np.inf, 1.0, -2 * np.pi, -np.inf],
                [np.inf, 100000.0, 2 * np.pi, np.inf]
            )
        )
        fitted = model(positions, *popt)
        residual = np.mean((signal - fitted) ** 2)
        variance = np.var(signal)
        r2 = 1.0 - residual / (variance + 1e-10) if variance > 1e-10 else 1.0
        return popt, model, r2
    except (RuntimeError, ValueError):
        return None, model, 0.0


# ================================================================
# Extension methods
# ================================================================

def quintic_step(t):
    t = np.clip(t, 0.0, 1.0)
    return t * t * t * (t * (t * 6 - 15) + 10)


def preserve_and_blend(result, original, blend_width=12):
    orig_len = original.shape[0]
    if blend_width <= 0:
        result[:orig_len] = original
        return result

    blend_start = max(0, orig_len - blend_width)
    result[:blend_start] = original[:blend_start]

    for i in range(blend_start, orig_len):
        alpha = quintic_step((i - blend_start) / blend_width)
        result[i] = (1.0 - alpha) * original[i] + alpha * result[i]

    return result


def match_norms(result, original, start_idx):
    orig_norms = np.linalg.norm(original, axis=1)
    mean_norm = np.mean(orig_norms)

    for i in range(start_idx, result.shape[0]):
        norm = np.linalg.norm(result[i])
        if norm > 1e-10:
            result[i] *= mean_norm / norm

    return result


def extend_learned_sinusoidal(pos_emb, target_length, n_components=3,
                               blend_width=12):
    """Fit N sinusoids + trend to each dimension."""
    orig_len, dim = pos_emb.shape
    pos_float = pos_emb.float().numpy().astype(np.float64)
    result = np.zeros((target_length, dim), dtype=np.float64)

    positions_orig = np.arange(orig_len, dtype=np.float64) / orig_len
    positions_ext = np.arange(target_length, dtype=np.float64) / orig_len

    fit_qualities = []
    t_start = time.time()

    for d in range(dim):
        if d % 50 == 0 or d == dim - 1:
            print_progress(d + 1, dim, "Fitting sinusoids")

        signal = pos_float[:, d]
        params, model_func, r2 = fit_sinusoidal_fast(signal, positions_orig, n_components)
        fit_qualities.append(r2)

        if params is not None and r2 > 0.3:
            result[:, d] = model_func(positions_ext, *params)
        else:
            result[:, d] = spectral_extrapolate_single(signal, target_length)

    elapsed = time.time() - t_start
    print(f"\r  Fitting sinusoids: done in {elapsed:.1f}s" + " " * 30)

    result = preserve_and_blend(result, pos_float, blend_width)
    result = match_norms(result, pos_float, orig_len)

    fq = np.array(fit_qualities)
    print(f"    Fit quality (R²): mean={np.mean(fq):.4f}, "
          f"min={np.min(fq):.4f}, max={np.max(fq):.4f}")
    print(f"    Good fits (R²>0.9):  {np.sum(fq > 0.9)}/{dim}")
    print(f"    Fair fits (R²>0.5):  {np.sum(fq > 0.5)}/{dim}")
    print(f"    Poor fits (R²<0.3):  {np.sum(fq < 0.3)}/{dim}")

    return torch.from_numpy(result).to(pos_emb.dtype)


def extend_vaswani_fitted(pos_emb, target_length, blend_width=12):
    """Fit Vaswani PE formula to each dimension."""
    orig_len, dim = pos_emb.shape
    pos_float = pos_emb.float().numpy().astype(np.float64)
    result = np.zeros((target_length, dim), dtype=np.float64)

    positions_orig = np.arange(orig_len, dtype=np.float64)
    positions_ext = np.arange(target_length, dtype=np.float64)

    good_fits = 0
    fallback_count = 0
    t_start = time.time()

    for d in range(dim):
        if d % 50 == 0 or d == dim - 1:
            print_progress(d + 1, dim, "Fitting Vaswani")

        signal = pos_float[:, d]
        params, model, r2 = fit_vaswani_fast(signal, d, dim, positions_orig)

        if params is not None and r2 > 0.5:
            result[:, d] = model(positions_ext, *params)
            good_fits += 1
        else:
            result[:, d] = spectral_extrapolate_single(signal, target_length)
            fallback_count += 1

    elapsed = time.time() - t_start
    print(f"\r  Fitting Vaswani: done in {elapsed:.1f}s" + " " * 30)
    print(f"    Vaswani fits: {good_fits}/{dim} good, {fallback_count}/{dim} spectral fallback")

    result = preserve_and_blend(result, pos_float, blend_width)
    result = match_norms(result, pos_float, orig_len)
    return torch.from_numpy(result).to(pos_emb.dtype)


def extend_hybrid(pos_emb, target_length, blend_width=12):
    """
    Try Vaswani first, then sinusoidal, then spectral fallback.
    Pick best R² for each dimension.
    """
    orig_len, dim = pos_emb.shape
    pos_float = pos_emb.float().numpy().astype(np.float64)
    result = np.zeros((target_length, dim), dtype=np.float64)

    positions_orig_norm = np.arange(orig_len, dtype=np.float64) / orig_len
    positions_ext_norm = np.arange(target_length, dtype=np.float64) / orig_len
    positions_orig_abs = np.arange(orig_len, dtype=np.float64)
    positions_ext_abs = np.arange(target_length, dtype=np.float64)

    method_counts = {"vaswani": 0, "sinusoidal": 0, "spectral": 0}
    r2_threshold = 0.7
    t_start = time.time()

    for d in range(dim):
        if d % 50 == 0 or d == dim - 1:
            elapsed = time.time() - t_start
            eta = (elapsed / (d + 1)) * (dim - d - 1) if d > 0 else 0
            print_progress(d + 1, dim, f"Hybrid (ETA {eta:.0f}s)")

        signal = pos_float[:, d]

        best_r2 = -1.0
        best_result = None
        best_method = "spectral"

        # Try Vaswani
        params_v, model_v, r2_v = fit_vaswani_fast(signal, d, dim, positions_orig_abs)
        if params_v is not None and r2_v > best_r2:
            best_r2 = r2_v
            best_result = model_v(positions_ext_abs, *params_v)
            best_method = "vaswani"

        # Try sinusoidal (only if Vaswani didn't reach threshold)
        if best_r2 < r2_threshold:
            params_s, model_s, r2_s = fit_sinusoidal_fast(
                signal, positions_orig_norm, n_components=3
            )
            if params_s is not None and r2_s > best_r2:
                best_r2 = r2_s
                best_result = model_s(positions_ext_norm, *params_s)
                best_method = "sinusoidal"

        # Spectral fallback if both poor
        if best_r2 < r2_threshold:
            best_result = spectral_extrapolate_single(signal, target_length)
            best_method = "spectral"

        result[:, d] = best_result
        method_counts[best_method] += 1

    elapsed = time.time() - t_start
    print(f"\r  Hybrid fitting: done in {elapsed:.1f}s" + " " * 40)
    print(f"    Method selection: {method_counts}")
    total = sum(method_counts.values())
    for m, c in method_counts.items():
        print(f"      {m}: {c}/{total} ({c/total*100:.1f}%)")

    result = preserve_and_blend(result, pos_float, blend_width)
    result = match_norms(result, pos_float, orig_len)
    return torch.from_numpy(result).to(pos_emb.dtype)


# ================================================================
# Analysis
# ================================================================

def analyze_embedding(name, original, extended, method_name):
    orig_len = original.shape[0]
    orig_np = original.float().numpy()
    ext_np = extended.float().numpy()

    diff = np.abs(orig_np - ext_np[:orig_len])

    print(f"\n  [{method_name}] {name}")
    print(f"    Shape: {tuple(original.shape)} -> {tuple(extended.shape)}")
    print(f"    Original preserved: max_diff={np.max(diff):.2e}, mean_diff={np.mean(diff):.2e}")

    orig_norms = np.linalg.norm(orig_np, axis=1)
    ext_norms = np.linalg.norm(ext_np[orig_len:], axis=1)

    print(f"    Original norms:  mean={np.mean(orig_norms):.4f} "
          f"std={np.std(orig_norms):.4f} "
          f"[{np.min(orig_norms):.4f}, {np.max(orig_norms):.4f}]")
    print(f"    Extended norms:  mean={np.mean(ext_norms):.4f} "
          f"std={np.std(ext_norms):.4f} "
          f"[{np.min(ext_norms):.4f}, {np.max(ext_norms):.4f}]")

    ratio = np.mean(ext_norms) / np.mean(orig_norms)
    quality = ("EXCELLENT" if 0.95 <= ratio <= 1.05 else
               "GOOD" if 0.9 <= ratio <= 1.1 else
               "ACCEPTABLE" if 0.8 <= ratio <= 1.2 else "WARNING")
    print(f"    Norm ratio: {ratio:.4f} — {quality}")

    last_orig = orig_np[-1]
    first_ext = ext_np[orig_len]
    cos_sim = np.dot(last_orig, first_ext) / (
        np.linalg.norm(last_orig) * np.linalg.norm(first_ext) + 1e-10)
    print(f"    Boundary cos_sim: {cos_sim:.6f}")

    inter_dists = [np.linalg.norm(orig_np[i + 1] - orig_np[i])
                   for i in range(orig_len - 1)]
    boundary_dist = np.linalg.norm(last_orig - first_ext)
    print(f"    Boundary/avg_step: {boundary_dist / (np.mean(inter_dists) + 1e-10):.4f}")


# ================================================================
# Main
# ================================================================

ALL_METHODS = ["learned_sinusoidal", "vaswani_fitted", "hybrid"]


def main():
    parser = argparse.ArgumentParser(
        description="SDXL Long Context — Learned Sinusoidal Extrapolation",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Methods:
  learned_sinusoidal  Fit N sinusoids + trend to each dimension
  vaswani_fitted      Fit original Transformer PE formula
  hybrid              Best of all methods per dimension (recommended)
        """
    )
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--target_len", type=int, default=MAX_LENGTH)
    parser.add_argument("--method", choices=ALL_METHODS, default="hybrid")
    parser.add_argument("--n_components", type=int, default=3)
    parser.add_argument("--blend_width", type=int, default=12)
    parser.add_argument("--analyze", action="store_true")

    args = parser.parse_args()
    target = min(args.target_len, MAX_LENGTH)

    print(f"SDXL Long Context — Learned Sinusoidal Extrapolation")
    print(f"=" * 55)
    print(f"Input:        {args.input}")
    print(f"Output:       {args.output}")
    print(f"Target:       {target} tokens")
    print(f"Method:       {args.method}")
    if args.method == "learned_sinusoidal":
        print(f"Components:   {args.n_components}")
    print(f"Blend width:  {args.blend_width}")

    print(f"\nLoading...")
    state_dict = {}
    with safe_open(args.input, framework="pt", device="cpu") as f:
        for key in f.keys():
            state_dict[key] = f.get_tensor(key)
    print(f"  {len(state_dict)} tensors")

    modified = []
    for key in POS_EMB_KEYS:
        if key not in state_dict:
            continue
        original = state_dict[key]
        if original.dim() != 2 or original.shape[0] >= target:
            continue

        print(f"\nExtending: {key}")
        print(f"  {tuple(original.shape)} -> ({target}, {original.shape[1]})")

        t0 = time.time()

        if args.method == "learned_sinusoidal":
            extended = extend_learned_sinusoidal(
                original, target, args.n_components, args.blend_width)
        elif args.method == "vaswani_fitted":
            extended = extend_vaswani_fitted(original, target, args.blend_width)
        elif args.method == "hybrid":
            extended = extend_hybrid(original, target, args.blend_width)

        elapsed = time.time() - t0
        print(f"  Total time: {elapsed:.1f}s")

        if args.analyze:
            analyze_embedding(key, original, extended, args.method)

        state_dict[key] = extended
        modified.append(key)

    if not modified:
        print("\nERROR: No positional embeddings found!")
        for key in state_dict:
            if "position" in key.lower():
                print(f"  {key}: {tuple(state_dict[key].shape)}")
        sys.exit(1)

    print(f"\nSaving: {args.output}")
    metadata = {
        "long_context": f"learned_{args.method}",
        "original_length": "77",
        "target_length": str(target),
        "method": args.method,
        "modified_keys": json.dumps(modified),
    }
    clean = {k: v.contiguous() for k, v in state_dict.items()}
    save_file(clean, args.output, metadata=metadata)

    size = os.path.getsize(args.output) / (1024 ** 3)
    print(f"\nDone! ({size:.2f} GB)")
    print(f"Modified: {len(modified)} tensors")
    for key in modified:
        print(f"  - {key}")


if __name__ == "__main__":
    main()