#!/usr/bin/env python3
"""
UV-Vis Spectrum Simulator v2 — optimum in the INTERIOR of parameter space.

The peak absorbance is maximized at ~55°C and ~pH 8.5 (not at the corners),
making acquisition function plots more visually informative for tutorials.

Physics model:
  - Peak height has a Gaussian-shaped optimum in (T, pH) space:
      height = 1.4 * exp(-0.5*((T-55)/15)^2 - 0.5*((pH-8.5)/2)^2) + 0.15
    This creates a smooth "hill" centered at T=55, pH=8.5.
  - Peak center: 450 + 0.3*(T-25) - 5*(pH-7) nm  (same red/blue shifts)
  - Peak width:  30 + 0.15*(T-25) nm

Usage:
  python simulate_spectra_v2.py init --output_dir ./spectra --n 8
  python simulate_spectra_v2.py run  --output_dir ./spectra \
      --params '{"temperature_C": 55, "pH": 8.5}'
"""

import argparse
import json
import sys
import numpy as np
import pandas as pd
from pathlib import Path


WAVELENGTHS = np.linspace(350, 600, 251)

BOUNDS = {
    "temperature_C": (5.0, 100.0),
    "pH": (1.0, 14.0),
}

# Ground truth optimum (interior of parameter space)
OPT_TEMP = 55.0
OPT_PH = 8.5


def generate_spectrum(temp: float, ph: float, noise_level: float = 0.005) -> np.ndarray:
    """Generate a synthetic UV-Vis spectrum with an interior optimum."""
    peak_center = 450 + 0.3 * (temp - 25) - 5.0 * (ph - 7.0)
    peak_width = 30 + 0.15 * (temp - 25)

    # Peak height: Gaussian hill centered at (OPT_TEMP, OPT_PH)
    peak_height = (
        1.4 * np.exp(
            -0.5 * ((temp - OPT_TEMP) / 15) ** 2
            - 0.5 * ((ph - OPT_PH) / 2) ** 2
        )
        + 0.15  # baseline floor
    )

    absorbance = peak_height * np.exp(
        -0.5 * ((WAVELENGTHS - peak_center) / peak_width) ** 2
    )
    absorbance += 0.15 * np.exp(-0.5 * ((WAVELENGTHS - 400) / 15) ** 2)
    baseline = 0.02 + 0.0001 * (WAVELENGTHS - 350)
    noise = np.random.normal(0, noise_level, len(WAVELENGTHS))

    return np.clip(absorbance + baseline + noise, 0, None)


def save_spectrum(temp: float, ph: float, output_dir: Path, label: str = None):
    absorbance = generate_spectrum(temp, ph)
    if label is None:
        label = f"spectrum_T{temp:.1f}_pH{ph:.1f}"
    label = label.replace(" ", "_")

    csv_path = output_dir / f"{label}.csv"
    json_path = output_dir / f"{label}.json"

    df = pd.DataFrame({"wavelength_nm": WAVELENGTHS, "absorbance": absorbance})
    df.to_csv(csv_path, index=False)

    conditions = {"temperature_C": round(temp, 2), "pH": round(ph, 2)}
    with open(json_path, "w") as f:
        json.dump(conditions, f, indent=2)

    return csv_path.name


def cmd_init(args):
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    n = args.n
    if args.grid:
        n_temp = max(2, int(np.sqrt(n)))
        n_ph = max(2, n // n_temp)
        temps = np.linspace(25, 85, n_temp)
        phs = np.linspace(4, 12, n_ph)
        conditions = [
            {"temperature_C": float(t), "pH": float(p)}
            for t in temps for p in phs
        ][:n]
    else:
        np.random.seed(args.seed)
        conditions = [
            {
                "temperature_C": round(np.random.uniform(20, 90), 1),
                "pH": round(np.random.uniform(3, 12), 1),
            }
            for _ in range(n)
        ]

    global_conditions = {}
    for cond in conditions:
        fname = save_spectrum(cond["temperature_C"], cond["pH"], output_dir)
        global_conditions[fname] = cond
        print(f"  Generated: {fname}")

    with open(output_dir / "conditions.json", "w") as f:
        json.dump(global_conditions, f, indent=2)

    print(f"\nGenerated {len(conditions)} initial spectra in {output_dir}/")
    print(f"Ground truth optimum: T={OPT_TEMP}C, pH={OPT_PH}")


def cmd_run(args):
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    try:
        params = json.loads(args.params)
    except json.JSONDecodeError as e:
        print(f"Error: could not parse --params JSON: {e}", file=sys.stderr)
        sys.exit(1)

    if isinstance(params, dict):
        params = [params]

    conditions_path = output_dir / "conditions.json"
    if conditions_path.exists():
        with open(conditions_path) as f:
            global_conditions = json.load(f)
    else:
        global_conditions = {}

    generated = []
    for i, p in enumerate(params):
        temp = p.get("temperature_C")
        ph = p.get("pH")
        if temp is None or ph is None:
            print(f"Error: entry {i} missing 'temperature_C' or 'pH': {p}", file=sys.stderr)
            continue

        temp = float(np.clip(temp, *BOUNDS["temperature_C"]))
        ph = float(np.clip(ph, *BOUNDS["pH"]))

        label = f"spectrum_T{temp:.1f}_pH{ph:.1f}"
        fname = save_spectrum(temp, ph, output_dir, label=label)
        global_conditions[fname] = {"temperature_C": round(temp, 2), "pH": round(ph, 2)}
        generated.append(fname)
        print(f"  Generated: {fname}  (T={temp:.1f}C, pH={ph:.1f})")

    with open(conditions_path, "w") as f:
        json.dump(global_conditions, f, indent=2)

    print(f"\nGenerated {len(generated)} new spectra in {output_dir}/")
    if generated:
        print(f"\nFiles to analyze:")
        for fname in generated:
            print(f"  {output_dir / fname}")


def main():
    parser = argparse.ArgumentParser(
        description="UV-Vis Spectrum Simulator v2 (interior optimum)"
    )
    sub = parser.add_subparsers(dest="command", required=True)

    p_init = sub.add_parser("init", help="Generate initial batch")
    p_init.add_argument("--output_dir", required=True)
    p_init.add_argument("--n", type=int, default=8)
    p_init.add_argument("--grid", action="store_true")
    p_init.add_argument("--seed", type=int, default=42)

    p_run = sub.add_parser("run", help="Generate spectra for BO-suggested params")
    p_run.add_argument("--output_dir", required=True)
    p_run.add_argument("--params", required=True)

    args = parser.parse_args()
    if args.command == "init":
        cmd_init(args)
    elif args.command == "run":
        cmd_run(args)


if __name__ == "__main__":
    main()
