import argparse
from pathlib import Path

import numpy as np
import yaml


def real_images(paths, size, scale_factor):
    import tifffile

    if not paths:
        raise FileNotFoundError("No TIFF files supplied")
    images = []
    scale_factors = []
    for path in paths:
        image = tifffile.imread(path)
        original_dtype = image.dtype
        if image.ndim != 3:
            raise ValueError(f"Expected a 13-band TIFF, got {image.shape} from {path}")
        if image.shape[0] == 13:
            image = image.transpose(1, 2, 0)
        if image.shape[-1] != 13:
            raise ValueError(f"Expected 13 Sentinel-2 bands, got {image.shape} from {path}")
        image = np.delete(image, 10, axis=-1)  # Exclude Sentinel-2 B10 as in the paper.
        y = np.linspace(0, image.shape[0] - 1, size).round().astype(int)
        x = np.linspace(0, image.shape[1] - 1, size).round().astype(int)
        image = image[y][:, x]
        factor = scale_factor
        if factor is None:
            minimum = float(np.nanmin(image))
            maximum = float(np.nanmax(image))
            already_normalized = minimum >= 0.0 and maximum <= 1.0
            factor = 10000.0 if not already_normalized and (
                np.issubdtype(original_dtype, np.integer) or maximum > 1.0
            ) else 1.0
        images.append(np.clip(image.astype(np.float32) / factor, 0, 1).transpose(2, 0, 1))
        scale_factors.append(factor)
    return np.asarray(images, dtype=np.float32), np.asarray(scale_factors, dtype=np.float32)


def synthetic_images(count, size, seed):
    rng = np.random.default_rng(seed)
    y, x = np.mgrid[0:size, 0:size].astype(np.float32) / max(size - 1, 1)
    images = []
    for index in range(count):
        phase = rng.uniform(0, 2 * np.pi)
        bands = []
        for band in range(12):
            pattern = 0.45 + 0.22 * np.sin((band + 1) * x + phase)
            pattern += 0.18 * np.cos((band / 3 + 1) * y - phase)
            pattern += rng.normal(0, 0.025, (size, size))
            bands.append(np.clip(pattern, 0, 1))
        images.append(bands)
    return np.asarray(images, dtype=np.float32)


def save_npz(output, images, source, protocol, normalization, scale_factors, stage):
    band_order = np.asarray(["B1", "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B9", "B11", "B12"])
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, images=images, data_source=np.asarray(source),
                        protocol=np.asarray(protocol), band_order=band_order,
                        normalization=np.asarray(normalization), scale_factors=scale_factors,
                        stage=np.asarray(stage))
    print(f"saved: {output} shape={images.shape} stage={stage} data_source={source}")


def main():
    parser = argparse.ArgumentParser(description="Generate compact 12-band spectral data")
    parser.add_argument("--config", default="conf/config.yaml")
    parser.add_argument("--real-dir", help="Convert official 13-band Sentinel-2 TIFF files")
    parser.add_argument("--scale-factor", default="auto",
                        help="TIFF divisor, or 'auto' (10000 for integer/range > 1; otherwise 1)")
    parser.add_argument("--stage", choices=("stage1", "stage2"), default="stage2",
                        help="Target stage for real TIFF conversion")
    args = parser.parse_args()
    with open(args.config, encoding="utf-8") as handle:
        config = yaml.safe_load(handle)
    if args.scale_factor == "auto":
        scale_factor = None
    else:
        scale_factor = float(args.scale_factor)
        if not np.isfinite(scale_factor) or scale_factor <= 0:
            raise ValueError("--scale-factor must be a positive finite number or 'auto'")
    if args.real_dir:
        stage = next(item for item in config["stages"] if item["name"] == args.stage)
        count = stage["train_samples"] + (config["data"]["test_samples"] if args.stage == "stage2" else 0)
        paths = sorted(Path(args.real_dir).rglob("*.tif"))
        if len(paths) < count:
            raise ValueError(f"Need at least {count} TIFF files, found {len(paths)}")
        images, scale_factors = real_images(paths[:count], stage["image_size"], scale_factor)
        normalization = "divide_by_scale_factor_then_clip_0_1"
        save_npz(Path(stage["train_path"]), images[:stage["train_samples"]], "real",
                 config["data"]["protocol"], normalization,
                 scale_factors[:stage["train_samples"]], args.stage)
        if args.stage == "stage2":
            save_npz(Path(config["data"]["test_path"]), images[stage["train_samples"]:], "real",
                     config["data"]["protocol"], normalization,
                     scale_factors[stage["train_samples"]:], "stage2")
    else:
        for index, stage in enumerate(config["stages"]):
            images = synthetic_images(stage["train_samples"], stage["image_size"], config["runtime"]["seed"] + index)
            save_npz(Path(stage["train_path"]), images, "synthetic", config["data"]["protocol"],
                     "already_0_1", np.ones(len(images), np.float32), stage["name"])
        test_size = config["stages"][-1]["image_size"]
        images = synthetic_images(config["data"]["test_samples"], test_size, config["runtime"]["seed"] + 2)
        save_npz(Path(config["data"]["test_path"]), images, "synthetic", config["data"]["protocol"],
                 "already_0_1", np.ones(len(images), np.float32), "stage2")


if __name__ == "__main__":
    main()