import math
import os
import threading
import tkinter as tk
from pathlib import Path
from tkinter import filedialog, messagebox, ttk
import time

import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file, save_file


# ============================================================
# Configuration
# ============================================================

PREDICTOR_PATH = "./model.safetensors"

HIDDEN_DIM = 384
NUM_LAYERS = 8

GLOBAL_FEATURES = 24
SHAPE_CONT_FEATURES = 32
MAX_SHAPE_DIMS = 8
BITS_PER_DIM = 16
SHAPE_BINARY_FEATURES = MAX_SHAPE_DIMS * BITS_PER_DIM

FEATURES = (
    8
    + 24
    + 64
    + GLOBAL_FEATURES
    + SHAPE_CONT_FEATURES
    + SHAPE_BINARY_FEATURES
)

CHUNK = 65536


# ============================================================
# Device
# ============================================================

def get_device():
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA GPU required.")
    return torch.device("cuda")


# ============================================================
# Shape helpers
# ============================================================

def requires_fp32(shape, limit):
    """
    Safety fallback.

    If any tensor dimension exceeds the configured limit,
    the tensor is processed entirely through the FP32 path
    (only relevant when the user selected FP64 mode + fallback).
    """
    return any(int(size) > limit for size in shape)


def is_simple_shape(shape):
    """True for scalars / pure 1-D (or effectively 1-D) tensors."""
    return sum(size > 1 for size in shape) <= 1


# ============================================================
# Coordinate helpers
# ============================================================

def gather_coordinates(tensor, coordinates):
    indices = tuple(
        coordinates[:, dim]
        for dim in range(coordinates.shape[1])
    )
    return tensor[indices]


def gather_neighbor(tensor, coordinates, dimension, offset):
    coords = coordinates.clone()

    coords[:, dimension] = (
        coords[:, dimension] + offset
    ).clamp(
        0,
        tensor.shape[dimension] - 1,
    )

    return gather_coordinates(tensor, coords)


def make_coordinates(linear, shape):
    coords = []
    rem = linear

    for dimension in reversed(range(len(shape))):
        size = shape[dimension]
        coords.append(rem % size)
        rem = rem // size

    coords.reverse()

    return torch.stack(coords, dim=1)


# ============================================================
# Feature generation
# ============================================================

def element_features(quantized, coordinates, dtype):
    """
    EXACTLY matches the training architecture.

    8 features:
      x
      abs(x)
      sign(x)
      sign(x) * log1p(abs(x))
      x^2
      x^3
      1 / (1 + abs(x))
      tanh(x)
    """

    x = gather_coordinates(
        quantized,
        coordinates,
    ).to(dtype)

    abs_x = x.abs()

    return torch.stack(
        [
            x,
            abs_x,
            torch.sign(x),
            torch.sign(x) * torch.log1p(abs_x),
            x * x,
            x * x * x,
            1.0 / (1.0 + abs_x),
            torch.tanh(x),
        ],
        dim=1,
    )


def local_features(tensor, coordinates, dtype):
    """
    EXACTLY matches the training architecture.

    Four dimensions maximum × eight statistics = 32,
    then truncated/padded to 24.
    """

    summaries = []

    for dim in range(min(tensor.ndim, 4)):
        values = torch.stack(
            [
                gather_neighbor(
                    tensor,
                    coordinates,
                    dim,
                    offset,
                ).to(dtype)
                for offset in (-2, -1, 0, 1, 2)
            ],
            dim=1,
        )

        center = values[:, 2]

        summaries.append(
            torch.stack(
                [
                    values.mean(1),
                    values.std(1),
                    values.min(1).values,
                    values.max(1).values,
                    center - values[:, 1],
                    values[:, 3] - center,
                    values[:, 3] - 2 * center + values[:, 1],
                    (
                        values[:, 4]
                        - 4 * values[:, 3]
                        + 6 * center
                        - 4 * values[:, 1]
                        + values[:, 0]
                    ),
                ],
                dim=1,
            )
        )

    if summaries:
        result = torch.cat(summaries, dim=1)

        if result.shape[1] >= 24:
            return result[:, :24]

        return F.pad(
            result,
            (0, 24 - result.shape[1]),
        )

    return torch.zeros(
        coordinates.shape[0],
        24,
        device=tensor.device,
        dtype=dtype,
    )


def dimension_features(shape, coordinates, dtype):
    """
    EXACTLY matches training.

    64 features.
    """

    count = coordinates.shape[0]

    result = torch.zeros(
        count,
        64,
        device=coordinates.device,
        dtype=dtype,
    )

    result[:, 0] = len(shape) / 8.0

    for dim in range(min(len(shape), 8)):
        size = max(shape[dim], 1)

        n = (
            coordinates[:, dim].to(dtype)
            / max(size - 1, 1)
        )

        base = 1 + dim * 7

        result[:, base:base + 7] = torch.stack(
            [
                n,
                torch.sin(n * math.pi),
                torch.cos(n * math.pi),
                torch.sin(n * 2 * math.pi),
                torch.cos(n * 2 * math.pi),
                n * n,
                n - 0.5,
            ],
            dim=1,
        )

    return result


def global_statistics(tensor, dtype):
    """
    EXACTLY matches the training script.

    IMPORTANT:
    This produces 23 values, not 24.
    """

    x = tensor.to(dtype).reshape(-1)

    if x.numel() > 65536:
        x = x[
            torch.randint(
                0,
                x.numel(),
                (65536,),
                device=x.device,
            )
        ]

    mean = x.mean()

    centered = x - mean

    std = torch.sqrt(
        centered.pow(2).mean() + 1e-12
    )

    normalized = centered / std

    return torch.stack(
        [
            mean,
            std,
            x.abs().mean(),
            torch.sqrt(
                x.pow(2).mean() + 1e-12
            ),
            x.min(),
            x.max(),

            *[
                torch.quantile(x, q)
                for q in (
                    0.01,
                    0.05,
                    0.10,
                    0.25,
                    0.50,
                    0.75,
                    0.90,
                    0.95,
                    0.99,
                )
            ],

            normalized.pow(3).mean(),
            normalized.pow(4).mean() - 3.0,

            (x > 0).to(dtype).mean(),
            (x < 0).to(dtype).mean(),
            (x == 0).to(dtype).mean(),

            (x - mean).abs().mean(),

            torch.sqrt(
                x.pow(2).mean() + 1e-12
            ),

            normalized.abs().mean(),

            torch.tanh(x).mean(),
        ]
    )


def shape_features_continuous(shape, count, device, dtype):
    """
    Original continuous 32-dimensional shape features.
    """

    result = torch.zeros(
        count,
        SHAPE_CONT_FEATURES,
        device=device,
        dtype=dtype,
    )

    total = math.prod(shape)

    result[:, 0] = len(shape) / 8.0
    result[:, 1] = math.log1p(float(total))

    for dim in range(min(len(shape), 8)):
        base = 2 + dim * 3
        size = float(shape[dim])

        result[:, base] = math.log1p(size)
        result[:, base + 1] = (
            size / max(float(total), 1.0)
        )
        result[:, base + 2] = float(size > 1)

    return result


def shape_features_binary(shape, count, device, dtype):
    """
    128 binary shape features.

    8 dimensions × 16 bits.
    """

    result = torch.zeros(
        count,
        SHAPE_BINARY_FEATURES,
        device=device,
        dtype=dtype,
    )

    for dim in range(
        min(len(shape), MAX_SHAPE_DIMS)
    ):
        size = int(shape[dim])
        base = dim * BITS_PER_DIM

        for bit in range(BITS_PER_DIM):
            result[:, base + bit] = float(
                (size >> bit) & 1
            )

    return result


def build_features(
    fp_high,
    quantized,
    coordinates,
    statistics,
    dtype,
):
    """
    EXACT feature layout from training.
    """

    element = element_features(
        quantized,
        coordinates,
        dtype,
    )

    local = local_features(
        quantized,
        coordinates,
        dtype,
    )

    dimension = dimension_features(
        tuple(fp_high.shape),
        coordinates,
        dtype,
    )

    statistics = (
        statistics
        .to(dtype)
        .unsqueeze(0)
        .expand(coordinates.shape[0], -1)
    )

    shape_cont = shape_features_continuous(
        tuple(fp_high.shape),
        coordinates.shape[0],
        coordinates.device,
        dtype,
    )

    shape_binary = shape_features_binary(
        tuple(fp_high.shape),
        coordinates.shape[0],
        coordinates.device,
        dtype,
    )

    features = torch.cat(
        [
            element,
            local,
            dimension,
            statistics,
            shape_cont,
            shape_binary,
        ],
        dim=1,
    )

    if features.shape[1] != FEATURES:
        raise RuntimeError(
            f"Feature mismatch: expected {FEATURES}, "
            f"got {features.shape[1]}"
        )

    return features


# ============================================================
# Model architecture
# ============================================================

class ResidualBlock(nn.Module):
    """
    EXACTLY matches the training script.
    """

    def __init__(self, hidden):
        super().__init__()

        self.norm = nn.LayerNorm(hidden)

        self.fc1 = nn.Linear(
            hidden,
            hidden * 2,
        )

        self.fc2 = nn.Linear(
            hidden * 2,
            hidden,
        )

        self.dropout = nn.Dropout(0.05)

    def forward(self, x):
        residual = x

        h = self.norm(x)
        h = F.gelu(self.fc1(h))
        h = self.dropout(self.fc2(h))

        return F.gelu(
            residual + h
        )


class UniversalWeightReconstructor(nn.Module):
    """
    EXACTLY matches the training architecture.

    Heads:

      element_head -> per-element residual

      tensor_head  -> scalar per-tensor correction

    Final reconstruction residual:

      element_pred + tensor_pred
    """

    def __init__(
        self,
        input_dim,
        hidden_dim,
        layers,
    ):
        super().__init__()

        self.input = nn.Sequential(
            nn.Linear(
                input_dim,
                hidden_dim,
            ),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
        )

        self.blocks = nn.ModuleList(
            [
                ResidualBlock(hidden_dim)
                for _ in range(layers)
            ]
        )

        self.element_head = nn.Sequential(
            nn.LayerNorm(hidden_dim),
            nn.Linear(
                hidden_dim,
                hidden_dim // 2,
            ),
            nn.GELU(),
            nn.Linear(
                hidden_dim // 2,
                1,
            ),
        )

        self.tensor_head = nn.Sequential(
            nn.LayerNorm(hidden_dim),
            nn.Linear(
                hidden_dim,
                hidden_dim // 4,
            ),
            nn.GELU(),
            nn.Linear(
                hidden_dim // 4,
                1,
            ),
        )

    def forward(self, x):
        h = self.input(x)

        for block in self.blocks:
            h = block(h)

        element_pred = (
            self.element_head(h)
            .squeeze(-1)
        )

        tensor_pred = (
            self.tensor_head(
                h.mean(
                    dim=0,
                    keepdim=True,
                )
            )
            .squeeze()
        )

        return (
            element_pred,
            tensor_pred,
        )


# ============================================================
# Predictor loading
# ============================================================

def load_predictor(device, model_dtype):
    predictor_path = Path(
        PREDICTOR_PATH
    )

    if not predictor_path.exists():
        raise FileNotFoundError(
            f"Predictor checkpoint not found:\n"
            f"{predictor_path}"
        )

    state = load_file(
        str(predictor_path),
        device="cpu",
    )

    model = UniversalWeightReconstructor(
        input_dim=FEATURES,
        hidden_dim=HIDDEN_DIM,
        layers=NUM_LAYERS,
    )

    expected_keys = set(
        model.state_dict().keys()
    )

    actual_keys = set(
        state.keys()
    )

    missing = expected_keys - actual_keys
    unexpected = actual_keys - expected_keys

    if missing or unexpected:
        raise RuntimeError(
            "Predictor checkpoint architecture mismatch.\n\n"
            f"Expected {FEATURES} input features.\n"
            f"Expected HIDDEN_DIM={HIDDEN_DIM}.\n"
            f"Expected NUM_LAYERS={NUM_LAYERS}.\n\n"
            f"Missing keys:\n{sorted(missing)}\n\n"
            f"Unexpected keys:\n{sorted(unexpected)}"
        )

    # Training checkpoint is FP64; cast to requested dtype.
    state = {
        key: value.to(model_dtype)
        for key, value in state.items()
    }

    model.load_state_dict(
        state,
        strict=True,
    )

    model = model.to(
        device=device,
        dtype=model_dtype,
    )

    model.eval()

    return model


# ============================================================
# Tensor prediction
# ============================================================

@torch.inference_mode()
def predict_tensor(
    predictor_fp64,
    predictor_fp32,
    fp_high,
    quantized,
    use_fp32_path,
    output_dtype,
):
    """
    Reconstruct one tensor.

    Normal (FP64 mode, small tensors):
        FP64 feature generation + FP64 predictor

    Safety fallback / pure FP32 mode:
        FP32 feature generation + FP32 predictor

    Final output dtype is controlled by the GUI selection
    (FP64 or FP32).
    """

    device = fp_high.device
    shape = tuple(fp_high.shape)
    n = fp_high.numel()

    if use_fp32_path:
        predictor = predictor_fp32
        compute_dtype = torch.float32
    else:
        predictor = predictor_fp64
        compute_dtype = torch.float64

    # Keep the source quantized representation as the base.
    output = torch.empty(
        n,
        device=device,
        dtype=compute_dtype,
    )

    # Global statistics must use the same compute precision
    # as the selected predictor path.
    statistics = global_statistics(
        fp_high,
        compute_dtype,
    )

    done = 0

    while done < n:
        end = min(
            done + CHUNK,
            n,
        )

        linear = torch.arange(
            done,
            end,
            device=device,
            dtype=torch.long,
        )

        coordinates = make_coordinates(
            linear,
            shape,
        )

        features = build_features(
            fp_high,
            quantized,
            coordinates,
            statistics,
            compute_dtype,
        )

        element_pred, tensor_pred = predictor(
            features
        )

        # EXACTLY matches training:
        #
        # combined_pred = element_pred + tensor_pred
        residual = (
            element_pred
            + tensor_pred
        )

        base = gather_coordinates(
            quantized,
            coordinates,
        ).to(compute_dtype)

        output[
            done:end
        ] = base + residual

        done = end

    return (
        output
        .to(output_dtype)
        .reshape(shape)
    )


# ============================================================
# Conversion
# ============================================================

def run_conversion(
    path,
    log,
    log_bold,
    save_path,
    skip_simple,
    prefer_fp64,
    use_large_fp32_fallback,
    fp32_shape_limit,
):
    device = get_device()

    output_dtype = torch.float64 if prefer_fp64 else torch.float32
    output_label = "FP64" if prefer_fp64 else "FP32"

    log(
        f"GPU: "
        f"{torch.cuda.get_device_name(0)}"
    )

    properties = (
        torch.cuda.get_device_properties(0)
    )

    log(
        f"VRAM: "
        f"{properties.total_memory / 1024**3:.2f} GB"
    )

    log(
        f"Predictor architecture: "
        f"{FEATURES} -> {HIDDEN_DIM} "
        f"with {NUM_LAYERS} residual blocks"
    )

    model_dtype = torch.float64 if prefer_fp64 else torch.float32

    log(
        f"Selected precision mode: "
        f"{output_label}"
    )

    if prefer_fp64 and use_large_fp32_fallback:
        log(
            f"Large-tensor FP32 fallback enabled: "
            f"any dim > {fp32_shape_limit:,}"
        )
    else:
        log(
            "Large-tensor FP32 fallback: disabled"
        )

    log(
        f"Output dtype: "
        f"{output_dtype}"
    )

    log(
        f"Skip 1D/simple: "
        f"{skip_simple}"
    )

    log(
        f"Loading source model: {path}"
    )

    state = load_file(
        path,
        device="cuda",
    )

    log(
        f"Loaded {len(state):,} tensors"
    )

    # Primary predictor (FP64 or FP32 according to mode).
    predictor_primary = load_predictor(
        device,
        model_dtype,
    )

    # Always keep an FP32 copy available for the safety path
    # (or pure FP32 mode).
    if prefer_fp64:
        predictor_fp64 = predictor_primary

        predictor_fp32 = (
            UniversalWeightReconstructor(
                input_dim=FEATURES,
                hidden_dim=HIDDEN_DIM,
                layers=NUM_LAYERS,
            )
            .to(
                device=device,
                dtype=torch.float32,
            )
        )

        predictor_fp32.load_state_dict(
            {
                key: value.to(torch.float32)
                for key, value in predictor_fp64.state_dict().items()
            },
            strict=True,
        )
        predictor_fp32.eval()

        log("Loaded matched FP64 predictor.")
        log("Created FP32 safety predictor.")
    else:
        predictor_fp64 = None
        predictor_fp32 = predictor_primary
        log("Loaded FP32 predictor.")

    output = {}

    total_tensors = len(state)

    converted = 0
    skipped = 0
    fp32_fallbacks = 0
    unsupported_casts = 0

    start_time = time.time()

    for index, (name, tensor) in enumerate(
        state.items(),
        1,
    ):
        shape = tuple(
            tensor.shape
        )

        # ----------------------------------------------------
        # Non-floating tensors
        # ----------------------------------------------------

        if not tensor.is_floating_point():
            output[name] = tensor

            log(
                f"[{index}/{total_tensors}] "
                f"{name} {shape} "
                f"unchanged (non-float)"
            )

            continue

        # ----------------------------------------------------
        # Source dtype handling
        # ----------------------------------------------------

        original_dtype = tensor.dtype
        src_label = str(original_dtype).replace("torch.", "")

        if original_dtype in (
            torch.float16,
            torch.bfloat16,
        ):
            # Accepted as-is (model was trained on these)
            quantized = tensor
            cast_note = ""
        else:
            # Accept other floating dtypes but warn + cast to BF16
            unsupported_casts += 1
            log_bold(
                f"WARNING: Tensor '{name}' has dtype {src_label}. "
                f"The model was NOT trained for this dtype. "
                f"Casting to BF16 before reconstruction."
            )
            quantized = tensor.to(torch.bfloat16)
            cast_note = f" (cast {src_label}→BF16)"

        # ----------------------------------------------------
        # Optional simple-shape bypass
        # ----------------------------------------------------

        if (
            skip_simple
            and is_simple_shape(shape)
        ):
            output[name] = (
                quantized
                .to(output_dtype)
                .contiguous()
            )

            skipped += 1

            log(
                f"[{index}/{total_tensors}] "
                f"{name} {shape} "
                f"simple shape -> {output_label} "
                f"(skipped reconstruction){cast_note}"
            )

            continue

        # ----------------------------------------------------
        # Precision selection
        # ----------------------------------------------------

        use_fp32_path = False

        if not prefer_fp64:
            # Pure FP32 mode
            use_fp32_path = True
            precision_label = "FP32"
        else:
            # FP64 mode – optional large-tensor fallback
            if (
                use_large_fp32_fallback
                and requires_fp32(shape, fp32_shape_limit)
            ):
                use_fp32_path = True
                fp32_fallbacks += 1
                precision_label = "FP32 SAFETY"
            else:
                use_fp32_path = False
                precision_label = "FP64"

        # High precision view used for global statistics and
        # shape/context processing.
        compute_dtype = (
            torch.float32 if use_fp32_path else torch.float64
        )
        fp_high = quantized.to(compute_dtype)

        try:
            predicted = predict_tensor(
                predictor_fp64,
                predictor_fp32,
                fp_high,
                quantized,
                use_fp32_path,
                output_dtype,
            )

            output[name] = (
                predicted
                .contiguous()
            )

            converted += 1

            if original_dtype == torch.bfloat16:
                src_dtype = "BF16"
            elif original_dtype == torch.float16:
                src_dtype = "FP16"
            else:
                src_dtype = src_label

            log(
                f"[{index}/{total_tensors}] "
                f"{name} {shape} "
                f"{src_dtype} -> "
                f"{precision_label} predictor -> "
                f"{output_label} reconstructed{cast_note}"
            )

        except Exception as exc:
            raise RuntimeError(
                f"\nReconstruction failed.\n"
                f"Tensor: {name}\n"
                f"Shape: {shape}\n"
                f"Source dtype: {original_dtype}\n"
                f"Predictor dtype: {compute_dtype}\n"
                f"Original error: {exc}"
            ) from exc

        finally:
            del quantized
            del fp_high

            if "predicted" in locals():
                del predicted

        if converted % 10 == 0:
            torch.cuda.empty_cache()

    # --------------------------------------------------------
    # Save
    # --------------------------------------------------------

    log("")
    log(
        f"Reconstructed tensors : "
        f"{converted:,}"
    )

    log(
        f"Simple tensors skipped: "
        f"{skipped:,}"
    )

    log(
        f"FP32 safety fallbacks : "
        f"{fp32_fallbacks:,}"
    )

    if unsupported_casts > 0:
        log_bold(
            f"Unsupported dtypes cast to BF16: "
            f"{unsupported_casts:,}"
        )

    log(
        f"Writing {output_label} model..."
    )

    save_file(
        output,
        save_path,
    )

    elapsed = (
        time.time()
        - start_time
    )

    log(
        f"Saved: {save_path}"
    )

    log(
        f"Time : {elapsed:.2f}s"
    )

    del output
    del state
    if predictor_fp64 is not None:
        del predictor_fp64
    del predictor_fp32

    torch.cuda.empty_cache()

    log("")
    log("DONE")


# ============================================================
# GUI
# ============================================================

class App:
    def __init__(self, root):
        self.root = root

        root.title(
            "Weight Reconstructor "
            "(FP16/BF16 → FP32/FP64)"
        )

        root.geometry(
            "1000x820"
        )

        self.path = tk.StringVar()

        # ---- Path row ----
        top = ttk.Frame(
            root,
            padding=10,
        )
        top.pack(fill="x")

        ttk.Label(
            top,
            text="Source Model:",
        ).pack(side="left")

        ttk.Entry(
            top,
            textvariable=self.path,
            width=70,
        ).pack(
            side="left",
            padx=8,
            fill="x",
            expand=True,
        )

        ttk.Button(
            top,
            text="Browse",
            command=self.browse,
        ).pack(side="left")

        # ---- Controls row 1 ----
        controls = ttk.Frame(
            root,
            padding=(10, 0),
        )
        controls.pack(fill="x")

        self.save = tk.BooleanVar(value=True)

        ttk.Checkbutton(
            controls,
            text="Save output",
            variable=self.save,
        ).pack(side="left")

        self.skip_simple = tk.BooleanVar(
            value=False          # default OFF
        )

        ttk.Checkbutton(
            controls,
            text=(
                "Skip 1D / simple-shape "
                "tensors (cast only)"
            ),
            variable=self.skip_simple,
        ).pack(
            side="left",
            padx=(20, 0),
        )

        self.convert_button = ttk.Button(
            controls,
            text="Convert",
            command=self.start,
        )
        self.convert_button.pack(
            side="left",
            padx=20,
        )

        # ---- Precision controls row ----
        precision_frame = ttk.LabelFrame(
            root,
            text="Precision mode (compute + save)",
            padding=10,
        )
        precision_frame.pack(
            fill="x",
            padx=10,
            pady=(8, 0),
        )

        self.precision_mode = tk.StringVar(value="FP64")

        ttk.Radiobutton(
            precision_frame,
            text="FP64 (higher precision)",
            variable=self.precision_mode,
            value="FP64",
            command=self._update_fallback_state,
        ).pack(side="left")

        ttk.Radiobutton(
            precision_frame,
            text="FP32 (faster / lower VRAM)",
            variable=self.precision_mode,
            value="FP32",
            command=self._update_fallback_state,
        ).pack(side="left", padx=(20, 0))

        # Large-tensor fallback controls
        self.use_large_fp32 = tk.BooleanVar(value=True)

        self.fallback_check = ttk.Checkbutton(
            precision_frame,
            text="Use FP32 for tensors larger than",
            variable=self.use_large_fp32,
            command=self._update_fallback_state,
        )
        self.fallback_check.pack(side="left", padx=(30, 4))

        self.fp32_limit = tk.StringVar(value="16326")

        self.limit_entry = ttk.Entry(
            precision_frame,
            textvariable=self.fp32_limit,
            width=8,
        )
        self.limit_entry.pack(side="left")

        ttk.Label(
            precision_frame,
            text="(per dimension)",
        ).pack(side="left", padx=(4, 0))

        # Initial state
        self._update_fallback_state()

        # ---- Log ----
        self.text = tk.Text(
            root,
            bg="#0d0d0d",
            fg="#dddddd",
            insertbackground="white",
            font=("Consolas", 9),
            wrap="none",
        )
        self.text.pack(
            fill="both",
            expand=True,
            padx=10,
            pady=10,
        )

        # Bold tag for warnings
        self.text.tag_configure(
            "bold",
            font=("Consolas", 9, "bold"),
            foreground="#ffcc00",
        )

    def _update_fallback_state(self):
        """Enable/disable the large-tensor controls based on mode."""
        is_fp64 = self.precision_mode.get() == "FP64"

        state = "normal" if is_fp64 else "disabled"
        self.fallback_check.configure(state=state)

        # Entry is only active when FP64 + checkbox checked
        if is_fp64 and self.use_large_fp32.get():
            self.limit_entry.configure(state="normal")
        else:
            self.limit_entry.configure(state="disabled")

    def browse(self):
        path = filedialog.askopenfilename(
            title=(
                "Select model "
                "(SafeTensors)"
            ),
            filetypes=[
                (
                    "SafeTensors",
                    "*.safetensors",
                ),
                (
                    "All files",
                    "*.*",
                ),
            ],
        )

        if path:
            self.path.set(path)

    def log(self, text):
        def write():
            self.text.insert(
                "end",
                text + "\n",
            )
            self.text.see("end")

        self.root.after(0, write)

    def log_bold(self, text):
        def write():
            self.text.insert(
                "end",
                text + "\n",
                "bold",
            )
            self.text.see("end")

        self.root.after(0, write)

    def start(self):
        path = self.path.get().strip()

        if not path:
            messagebox.showerror(
                "Error",
                "Select a source model first.",
            )
            return

        if not os.path.isfile(path):
            messagebox.showerror(
                "Error",
                "Selected model does not exist.",
            )
            return

        if (
            os.path.abspath(path)
            == os.path.abspath(PREDICTOR_PATH)
        ):
            messagebox.showerror(
                "Error",
                "The selected model cannot "
                "also be the predictor checkpoint.",
            )
            return

        # Validate the limit entry
        try:
            limit = int(self.fp32_limit.get().strip())
            if limit < 1:
                raise ValueError
        except ValueError:
            messagebox.showerror(
                "Error",
                "Tensor size limit must be a positive integer.",
            )
            return

        prefer_fp64 = self.precision_mode.get() == "FP64"
        output_label = "fp64" if prefer_fp64 else "fp32"

        base = os.path.splitext(path)[0]
        save_path = base + f"_predicted_{output_label}.safetensors"

        self.text.delete("1.0", "end")
        self.convert_button.configure(state="disabled")

        skip_simple = self.skip_simple.get()
        use_large_fp32 = self.use_large_fp32.get() and prefer_fp64

        self.log("Starting conversion...")
        self.log(
            "Architecture: "
            f"{FEATURES} features / "
            f"{HIDDEN_DIM} hidden / "
            f"{NUM_LAYERS} residual blocks"
        )
        self.log(
            f"Precision mode (compute + save): "
            f"{'FP64' if prefer_fp64 else 'FP32'}"
        )
        if prefer_fp64 and use_large_fp32:
            self.log(
                f"Large-tensor FP32 fallback: "
                f"enabled (limit {limit:,})"
            )
        else:
            self.log("Large-tensor FP32 fallback: disabled")
        self.log(
            f"Skip 1D/simple tensors: {skip_simple}"
        )

        threading.Thread(
            target=self.worker,
            args=(
                path,
                save_path,
                skip_simple,
                prefer_fp64,
                use_large_fp32,
                limit,
            ),
            daemon=True,
        ).start()

    def worker(
        self,
        path,
        save_path,
        skip_simple,
        prefer_fp64,
        use_large_fp32_fallback,
        fp32_shape_limit,
    ):
        try:
            run_conversion(
                path,
                self.log,
                self.log_bold,
                save_path,
                skip_simple,
                prefer_fp64,
                use_large_fp32_fallback,
                fp32_shape_limit,
            )

            output_label = "FP64" if prefer_fp64 else "FP32"

            self.root.after(
                0,
                lambda: (
                    self.convert_button.configure(
                        state="normal"
                    ),
                    messagebox.showinfo(
                        "Complete",
                        f"{output_label} reconstruction finished.",
                    ),
                ),
            )

        except Exception as exc:
            self.log("")
            self.log("ERROR:")
            self.log(str(exc))

            self.root.after(
                0,
                lambda: (
                    self.convert_button.configure(
                        state="normal"
                    ),
                    messagebox.showerror(
                        "Conversion Error",
                        str(exc),
                    ),
                ),
            )


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

if __name__ == "__main__":
    root = tk.Tk()
    App(root)
    root.mainloop()