from __future__ import annotations

from datetime import date, timedelta
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image

from .dataset import (
    DEFAULT_GROUP_FIELDS,
    assess_model_promotion,
    assign_splits,
    matrix_from_records,
    records_manifest,
    split_by_group,
    split_receipt,
)
from .features import extract_patch_features
from .inference import infer_capture
from .io import load_capture, save_reflectance, save_rgb, sha256_file, sha256_tree, write_json
from .model import build_model_card, evaluate_binary, fit_logistic, load_model, save_model
from .quality import quality_check_capture


SYNTHETIC_WARNING = "Deterministic synthetic engineering fixture; not field plant-disease evidence or accuracy."


def _synthetic_scene(seed: int, size: int = 128) -> tuple[np.ndarray, dict[str, np.ndarray], np.ndarray]:
    rng = np.random.default_rng(seed)
    rows, columns = np.mgrid[0:size, 0:size]
    texture = 0.02 * np.sin(columns / 3.7) + 0.015 * np.cos(rows / 4.9) + rng.normal(0, 0.008, (size, size))
    canopy = np.clip(0.72 + texture, 0.55, 0.88)
    offset_a = int(rng.integers(-8, 9))
    offset_b = int(rng.integers(-8, 9))
    lesion_a = ((rows - (34 + offset_a)) ** 2 / 18**2 + (columns - (45 + offset_b)) ** 2 / 22**2) <= 1
    lesion_b = ((rows - (91 - offset_b)) ** 2 / 20**2 + (columns - (86 + offset_a)) ** 2 / 17**2) <= 1
    streak = (rows > 52 + offset_a // 2) & (rows < 65 + offset_a // 2) & (columns > 12) & (columns < 54)
    lesion = lesion_a | lesion_b | streak
    confounder = ((rows - 28) ** 2 / 20**2 + (columns - 103) ** 2 / 17**2) <= 1
    shadow = (rows > 108) & (columns < 28)

    noise = rng.normal(0, 0.004, (size, size))
    bands = {
        "green": 0.18 + 0.05 * (1 - canopy) + 0.07 * lesion + 0.025 * confounder + noise,
        "red": 0.10 + 0.06 * (1 - canopy) + 0.18 * lesion + 0.06 * confounder + noise,
        "red_edge": 0.28 + 0.12 * canopy - 0.13 * lesion - 0.07 * confounder + noise,
        "nir": 0.48 + 0.34 * canopy - 0.31 * lesion - 0.17 * confounder + noise,
    }
    bands = {name: np.clip(value * np.where(shadow, 0.45, 1.0), 0.01, 0.95).astype(np.float32) for name, value in bands.items()}

    rgb = np.stack([
        0.18 + 0.30 * lesion + 0.12 * confounder + texture,
        0.54 - 0.25 * lesion - 0.11 * confounder + texture,
        0.15 - 0.045 * lesion - 0.02 * confounder + texture * 0.5,
    ], axis=2)
    rgb *= np.where(shadow[:, :, None], 0.48, 1.0)
    rgb = np.clip(rgb, 0.01, 0.96).astype(np.float32)
    return rgb, bands, lesion.astype(np.uint8)


def generate_synthetic_fixture(output_dir: str | Path, *, seed: int = 20260804, field_count: int = 6) -> list[Path]:
    if field_count < 3:
        raise ValueError("At least three synthetic fields are required")
    output = Path(output_dir)
    output.mkdir(parents=True, exist_ok=True)
    manifests: list[Path] = []
    start = date(2026, 8, 4)
    for index in range(field_count):
        field_id = f"SYN-F{index + 1:02d}"
        flight_id = f"{field_id}-FL01"
        capture_id = f"{flight_id}-CAP01"
        capture_date = (start + timedelta(days=index)).isoformat()
        capture_dir = output / capture_id
        rgb, bands, lesion = _synthetic_scene(seed + index * 1009)
        rgb_path = save_rgb(capture_dir / "rgb.png", rgb)
        band_paths = {name: save_reflectance(capture_dir / f"{name}.tif", value) for name, value in bands.items()}
        label_path = capture_dir / "latent-lesion-mask.png"
        Image.fromarray((lesion * 255).astype(np.uint8), mode="L").save(label_path)
        relative_files = {"rgb": rgb_path.name, **{name: path.name for name, path in band_paths.items()}}
        manifest = {
            "schemaVersion": 1,
            "dataProvenance": "synthetic",
            "containsFieldObservation": False,
            "captureId": capture_id,
            "fieldId": field_id,
            "flightId": flight_id,
            "captureDate": capture_date,
            "sensorModel": "DJI Mavic 3 Multispectral-compatible synthetic fixture",
            "scaleFactor": 10000,
            "processingLevel": "orthomosaic_reflectance_rgb_coregistered",
            "measurementScale": "relative_reflectance",
            "calibrationMethod": "synthetic_known_latent_scene",
            "crs": "LOCAL_SYNTHETIC_METRES",
            "gridTransform": [0.10, 0.0, 0.0, 0.0, -0.10, 12.8],
            "rgbToMultispectralRegistered": True,
            "coRegistrationRmsePx": 0.0,
            "bandMapValidated": True,
            "files": relative_files,
            "fileSha256": {name: sha256_file(capture_dir / relative) for name, relative in relative_files.items()},
            "labelMask": label_path.name,
            "labelReference": "synthetic_latent_truth",
            "warning": SYNTHETIC_WARNING,
        }
        manifest_path = write_json(capture_dir / "capture-manifest.json", manifest)
        manifests.append(manifest_path)
    write_json(output / "fixture-manifest.json", {
        "schemaVersion": 1,
        "seed": seed,
        "fieldCount": field_count,
        "captures": [str(path.relative_to(output).as_posix()) for path in manifests],
        "fieldPerformanceEstablished": False,
        "warning": SYNTHETIC_WARNING,
    })
    return manifests


def records_from_capture(manifest_path: str | Path, *, patch_size: int = 32, positive_fraction: float = 0.08) -> list[dict[str, Any]]:
    path = Path(manifest_path)
    rgb, bands, metadata = load_capture(path)
    qc = quality_check_capture(rgb, bands, metadata)
    if not qc["passed"]:
        raise RuntimeError("Synthetic capture unexpectedly failed QC: " + ", ".join(qc["failures"]))
    label_path = path.parent / str(metadata["labelMask"])
    label_mask = np.asarray(Image.open(label_path).convert("L"), dtype=np.float32) / 255.0
    if label_mask.shape != rgb.shape[:2]:
        raise ValueError("Label mask must share the registered image grid")
    records: list[dict[str, Any]] = []
    rows, columns = rgb.shape[:2]
    for top in range(0, rows - patch_size + 1, patch_size):
        for left in range(0, columns - patch_size + 1, patch_size):
            row_slice = slice(top, top + patch_size)
            column_slice = slice(left, left + patch_size)
            features = extract_patch_features(
                rgb[row_slice, column_slice, :],
                {name: value[row_slice, column_slice] for name, value in bands.items()},
            )
            fraction = float(label_mask[row_slice, column_slice].mean())
            sample_id = f"{metadata['captureId']}-R{top:04d}-C{left:04d}"
            records.append({
                "sample_id": sample_id,
                "field_id": metadata["fieldId"],
                "flight_id": metadata["flightId"],
                "capture_date": metadata["captureDate"],
                "data_provenance": "synthetic",
                "label": int(fraction >= positive_fraction),
                "label_reference": "synthetic_latent_truth",
                "latent_positive_fraction": round(fraction, 8),
                "features": features,
            })
    return records


def generate_records(manifests: list[Path], *, patch_size: int = 32) -> list[dict[str, Any]]:
    return [record for path in manifests for record in records_from_capture(path, patch_size=patch_size)]


def run_demo(output_dir: str | Path, *, seed: int = 20260804) -> dict[str, Any]:
    output = Path(output_dir)
    fixture_dir = output / "fixture"
    artifacts_dir = output / "artifacts"
    manifests = generate_synthetic_fixture(fixture_dir, seed=seed)
    records = generate_records(manifests)
    assignments = split_by_group(records, seed=seed, group_fields=DEFAULT_GROUP_FIELDS)
    assigned = assign_splits(records, assignments, group_fields=DEFAULT_GROUP_FIELDS)
    receipt = split_receipt(records, assignments, group_fields=DEFAULT_GROUP_FIELDS, seed=seed)
    train_x, train_y, _ = matrix_from_records(assigned, split="train")
    test_x, test_y, _ = matrix_from_records(assigned, split="test")
    model = fit_logistic(train_x, train_y)
    promotion = assess_model_promotion(assigned)
    if promotion["status"] != "RESEARCH_ONLY":
        raise AssertionError("Synthetic data must never promote a disease-screening model")
    model_path = save_model(artifacts_dir / "model.json", model)
    loaded = load_model(model_path)
    evaluation = evaluate_binary(loaded, test_x, test_y)
    card = build_model_card(loaded, split_receipt=receipt, evaluation=evaluation, promotion=promotion)
    write_json(artifacts_dir / "model-card.json", card)
    write_json(artifacts_dir / "dataset-records.json", {"schemaVersion": 1, "records": assigned})
    write_json(artifacts_dir / "dataset-manifest.json", records_manifest(records))
    write_json(artifacts_dir / "split-receipt.json", receipt)
    write_json(artifacts_dir / "evaluation.json", evaluation)
    inference = infer_capture(loaded, manifests[-1], patch_size=32, output_path=artifacts_dir / "inference-receipt.json")
    summary = {
        "schemaVersion": 1,
        "status": "PASSED",
        "dataProvenance": "synthetic",
        "fieldFlightExecuted": False,
        "fieldDataCollected": False,
        "fieldPerformanceEstablished": False,
        "modelPromotionStatus": promotion["status"],
        "syntheticFieldCount": len(manifests),
        "syntheticSampleCount": len(records),
        "testSampleCount": evaluation["sampleCount"],
        "inferenceStatus": inference["status"],
        "sourceTreeSha256": sha256_tree(Path(__file__).resolve().parent),
        "warning": SYNTHETIC_WARNING,
    }
    write_json(output / "engineering-demo-manifest.json", summary)
    return summary
