qdot/perception/dataset.py

"""
qdot/perception/dataset.py
==========================
CIMDataset — synthetic training data generator for the Inspection Agent.

KEY FIX (root cause of Phase 2 failure):
    _simulate() now centres each sample's scan window on the calculated
    charge transition voltage V_centre = -E_c_mean / lever_arm, spanning
    ±1.5 Coulomb periods.  The original fixed v_range = (-1.5, 1.5) V
    placed the transition entirely outside the window for all double-dot
    and single-dot samples (transition at -2.1 to -15.7 V), so the CNN
    learned gradient direction instead of charge morphology.

QFlow's role: transfer benchmark only — never in the training loop.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Optional, Tuple

import numpy as np

from qdot.core.types import ChargeLabel, Measurement, MeasurementModality, VoltagePoint


@dataclass
class DatasetConfig:
    n_per_class: int = 17_000
    resolutions: dict = field(default_factory=lambda: {16: 0.3, 32: 0.5, 64: 0.2})
    # v_range retained for reference; _simulate() ignores it in favour of
    # the per-sample transition-centred window.
    v_range: Tuple[float, float] = (-1.5, 1.5)
    seed: Optional[int] = 42
    augment: bool = True
    noise_aug_sigma: float = 0.02
    blur_aug_prob: float = 0.3
    # Double-dot
    dd_E_c_range: Tuple[float, float] = (1.0, 5.5)
    dd_t_c_range: Tuple[float, float] = (0.05, 0.6)
    dd_T_range: Tuple[float, float] = (0.01, 0.12)
    dd_lever_range: Tuple[float, float] = (0.35, 0.85)
    dd_asymmetry_max: float = 0.3
    # Single-dot
    sd_E_c_range: Tuple[float, float] = (1.5, 5.5)
    sd_t_c_range: Tuple[float, float] = (0.6, 1.5)
    sd_T_range: Tuple[float, float] = (0.01, 0.15)
    sd_lever_range: Tuple[float, float] = (0.3, 0.9)
    # Misc
    misc_E_c_range_sc: Tuple[float, float] = (0.3, 1.2)
    misc_E_c_range_barrier: Tuple[float, float] = (6.0, 12.0)
    misc_noise_range: Tuple[float, float] = (0.05, 0.15)


class CIMDataset:
    """
    Generates labelled 2D stability diagrams.

    Each sample's scan window is centred on V_centre = -E_c_mean / lever_arm
    with half-width δ = 1.5 × (one Coulomb period in voltage).  This
    guarantees honeycomb/diamond topology is visible in every training image.
    Uses the vectorised current_grid path (~100× faster than the old loop).
    """

    LABEL_MAP = {ChargeLabel.DOUBLE_DOT: 0, ChargeLabel.SINGLE_DOT: 1, ChargeLabel.MISC: 2}
    INT_TO_LABEL = {v: k for k, v in LABEL_MAP.items()}

    def __init__(self, config: Optional[DatasetConfig] = None) -> None:
        self.cfg = config or DatasetConfig()
        self.rng = np.random.default_rng(self.cfg.seed)

    def generate(self) -> Tuple[np.ndarray, np.ndarray]:
        samples = self.generate_measurements()
        arrays = np.stack([self._resize_to_64(s[0]) for s in samples], axis=0)
        arrays = arrays[:, np.newaxis, :, :]
        labels = np.array([self.LABEL_MAP[s[1]] for s in samples], dtype=np.int64)
        idx = self.rng.permutation(len(labels))
        return arrays[idx].astype(np.float32), labels[idx]

    def generate_measurements(self) -> List[Tuple[np.ndarray, ChargeLabel]]:
        samples: List[Tuple[np.ndarray, ChargeLabel]] = []
        n = self.cfg.n_per_class
        print(f"Generating {n} double-dot samples...")
        for _ in range(n):
            samples.append(self._generate_double_dot())
        print(f"Generating {n} single-dot samples...")
        for _ in range(n):
            samples.append(self._generate_single_dot())
        print(f"Generating {n} misc samples...")
        for _ in range(n):
            samples.append(self._generate_misc())
        print(f"Dataset complete: {len(samples)} samples.")
        return samples

    def _generate_double_dot(self) -> Tuple[np.ndarray, ChargeLabel]:
        cfg = self.cfg
        E_c_mean = self.rng.uniform(*cfg.dd_E_c_range)
        asym = self.rng.uniform(0, cfg.dd_asymmetry_max) * E_c_mean
        E_c1 = E_c_mean + asym / 2
        E_c2 = E_c_mean - asym / 2
        t_c = self.rng.uniform(*cfg.dd_t_c_range)
        T = self.rng.uniform(*cfg.dd_T_range)
        lever = self.rng.uniform(*cfg.dd_lever_range)
        noise = self.rng.uniform(0.005, 0.05)
        res = self._sample_resolution()
        arr = self._simulate(E_c1, E_c2, t_c, T, lever, noise, res)
        if cfg.augment:
            arr = self._augment(arr)
        return arr, ChargeLabel.DOUBLE_DOT

    def _generate_single_dot(self) -> Tuple[np.ndarray, ChargeLabel]:
        cfg = self.cfg
        mode = self.rng.choice(["strong_coupling", "asymmetric"])
        if mode == "strong_coupling":
            E_c1 = self.rng.uniform(*cfg.sd_E_c_range)
            E_c2 = self.rng.uniform(*cfg.sd_E_c_range)
            t_c = self.rng.uniform(*cfg.sd_t_c_range)
        else:
            E_c1 = self.rng.uniform(1.5, 4.0)
            E_c2 = E_c1 * self.rng.uniform(4.0, 8.0)
            t_c = self.rng.uniform(0.05, 0.4)
        T = self.rng.uniform(*cfg.sd_T_range)
        lever = self.rng.uniform(*cfg.sd_lever_range)
        noise = self.rng.uniform(0.005, 0.06)
        res = self._sample_resolution()
        arr = self._simulate(E_c1, E_c2, t_c, T, lever, noise, res)
        if cfg.augment:
            arr = self._augment(arr)
        return arr, ChargeLabel.SINGLE_DOT

    def _generate_misc(self) -> Tuple[np.ndarray, ChargeLabel]:
        cfg = self.cfg
        mode = self.rng.choice(["sc", "barrier", "high_noise"])
        if mode == "sc":
            E_c1 = self.rng.uniform(*cfg.misc_E_c_range_sc)
            E_c2 = self.rng.uniform(*cfg.misc_E_c_range_sc)
            t_c = self.rng.uniform(0.5, 2.0)
            T = self.rng.uniform(0.15, 0.5)
            lever = self.rng.uniform(0.2, 0.6)
            noise = self.rng.uniform(0.005, 0.04)
        elif mode == "barrier":
            E_c1 = self.rng.uniform(*cfg.misc_E_c_range_barrier)
            E_c2 = self.rng.uniform(*cfg.misc_E_c_range_barrier)
            t_c = self.rng.uniform(0.01, 0.1)
            T = self.rng.uniform(0.01, 0.06)
            lever = self.rng.uniform(0.1, 0.4)
            noise = self.rng.uniform(0.005, 0.04)
        else:
            E_c1 = self.rng.uniform(1.5, 5.0)
            E_c2 = self.rng.uniform(1.5, 5.0)
            t_c = self.rng.uniform(0.1, 0.5)
            T = self.rng.uniform(0.01, 0.1)
            lever = self.rng.uniform(0.3, 0.8)
            noise = self.rng.uniform(*cfg.misc_noise_range)
        res = self._sample_resolution()
        arr = self._simulate(E_c1, E_c2, t_c, T, lever, noise, res)
        if cfg.augment:
            arr = self._augment(arr)
        return arr, ChargeLabel.MISC

    def _simulate(
        self,
        E_c1: float, E_c2: float, t_c: float, T: float,
        lever: float, noise: float, res: int,
    ) -> np.ndarray:
        """
        Simulate a 2D stability diagram with a TRANSITION-CENTRED scan window.

        Window:  V_centre = -E_c_mean / lever  (charge degeneracy voltage)
                 half-width δ = 1.5 × Coulomb_period_in_voltage
                             = 1.5 × (E_c_mean / lever)

        This guarantees the honeycomb / Coulomb diamond features are present
        in every training image regardless of E_c and lever_arm values.
        Uses vectorised current_grid — no Python loop.
        """
        from qdot.simulator.cim import ConstantInteractionDevice

        device = ConstantInteractionDevice(
            E_c1=float(E_c1), E_c2=float(E_c2), t_c=float(t_c),
            T=float(T), lever_arm=float(lever), noise_level=float(noise),
            seed=int(self.rng.integers(0, 2**31)),
        )

        E_c_mean = (float(E_c1) + float(E_c2)) / 2.0
        V_centre = -E_c_mean / float(lever)
        coulomb_period_V = E_c_mean / float(lever)
        delta = max(1.5 * coulomb_period_V, 0.5)   # at least ±0.5 V

        v1_grid = np.linspace(V_centre - delta, V_centre + delta, res, dtype=np.float32)
        v2_grid = np.linspace(V_centre - delta, V_centre + delta, res, dtype=np.float32)

        patch = device.current_grid(v1_grid, v2_grid)

        lo, hi = patch.min(), patch.max()
        if hi - lo > 1e-12:
            patch = (patch - lo) / (hi - lo)
        else:
            patch = np.full_like(patch, 0.5)

        return patch.astype(np.float32)

    def _augment(self, arr: np.ndarray) -> np.ndarray:
        arr = arr.copy()
        sigma = self.rng.uniform(0, self.cfg.noise_aug_sigma)
        arr += self.rng.normal(0, sigma, arr.shape).astype(np.float32)
        k = int(self.rng.integers(0, 4))
        arr = np.rot90(arr, k=k)
        if self.rng.random() > 0.5:
            arr = np.fliplr(arr)
        if self.rng.random() > 0.5:
            arr = np.flipud(arr)
        if self.rng.random() < self.cfg.blur_aug_prob:
            from scipy.ndimage import gaussian_filter
            sigma_blur = self.rng.uniform(0.3, 1.2)
            arr = gaussian_filter(arr, sigma=sigma_blur).astype(np.float32)
        return np.clip(arr, 0.0, 1.0).astype(np.float32)

    def _sample_resolution(self) -> int:
        resolutions = list(self.cfg.resolutions.keys())
        weights = list(self.cfg.resolutions.values())
        total = sum(weights)
        probs = [w / total for w in weights]
        idx = self.rng.choice(len(resolutions), p=probs)
        return resolutions[idx]

    @staticmethod
    def _resize_to_64(arr: np.ndarray) -> np.ndarray:
        if arr.shape == (64, 64):
            return arr.astype(np.float32)
        from scipy.ndimage import zoom
        scale = 64.0 / arr.shape[0]
        resized = zoom(arr.astype(np.float64), scale, order=1)
        return np.clip(resized, 0.0, 1.0).astype(np.float32)

    @staticmethod
    def split(
        arrays: np.ndarray, labels: np.ndarray,
        val_frac: float = 0.15, seed: int = 42,
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        rng = np.random.default_rng(seed)
        classes = np.unique(labels)
        train_idx, val_idx = [], []
        for c in classes:
            idx = np.where(labels == c)[0]
            idx = rng.permutation(idx)
            n_val = max(1, int(len(idx) * val_frac))
            val_idx.extend(idx[:n_val].tolist())
            train_idx.extend(idx[n_val:].tolist())
        train_idx = np.array(train_idx)
        val_idx = np.array(val_idx)
        return arrays[train_idx], arrays[val_idx], labels[train_idx], labels[val_idx]