"""Train the plant-disease and plant-identification CNNs (transfer learning).

This produces the model files AgroSense's vision module loads. It needs TensorFlow
and the image datasets, so it is meant to be run OUTSIDE the lightweight app sandbox
(e.g. Colab / a GPU box) — not in the offline POC environment.

Datasets:
  * Disease  : PlantVillage (~54k images, 38 crop-disease classes). Widely mirrored
               on Kaggle (e.g. "PlantVillage Dataset"). Check the licence of the
               mirror you use before redistribution.
  * Plant ID : any labelled leaf/plant dataset arranged as one folder per class
               (e.g. Flavia leaf dataset, Oxford-102 Flowers, or an MIT-licensed
               Kaggle plant dataset). Folder layout: dataset/<class_name>/<images>.
  * Pest     : IP102 (~75k images, 102 crop-pest classes) or the Kaggle "Agricultural
               Pests Image Dataset" (12 classes). Same per-class-folder layout.

Usage:
    pip install "tensorflow>=2.15"
    python scripts/train_plant_models.py --data path/to/plantvillage --task disease \
        --out models/disease.keras --labels models/disease_labels.txt
    python scripts/train_plant_models.py --data path/to/plantid --task plant \
        --out models/plant.keras --labels models/plant_labels.txt
    python scripts/train_plant_models.py --data path/to/IP102 --task pest \
        --out models/pest.keras --labels models/pest_labels.txt

Then point AgroSense at them:
    AGROSENSE_VISION_DISEASE_MODEL=models/disease.keras
    AGROSENSE_VISION_DISEASE_LABELS=models/disease_labels.txt
    AGROSENSE_VISION_PLANT_MODEL=models/plant.keras
    AGROSENSE_VISION_PLANT_LABELS=models/plant_labels.txt

The model bakes in resizing+rescaling so the app's inference can feed raw [0,255]
224x224 RGB arrays (matching agrosense/vision.py KerasVisionModel).
"""
from __future__ import annotations

import argparse
from pathlib import Path

IMG_SIZE = (224, 224)


def build_model(num_classes: int):
    import tensorflow as tf
    from tensorflow.keras import layers, models

    base = tf.keras.applications.MobileNetV2(
        input_shape=IMG_SIZE + (3,), include_top=False, weights="imagenet")
    base.trainable = False
    return models.Sequential([
        layers.Input(shape=IMG_SIZE + (3,)),
        layers.Rescaling(1.0 / 127.5, offset=-1.0),   # MobileNetV2 expects [-1,1]
        base,
        layers.GlobalAveragePooling2D(),
        layers.Dropout(0.2),
        layers.Dense(num_classes, activation="softmax"),
    ])


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True, help="dataset dir (one subfolder per class)")
    ap.add_argument("--task", choices=["disease", "plant", "pest"], required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--labels", required=True)
    ap.add_argument("--epochs", type=int, default=8)
    ap.add_argument("--batch", type=int, default=32)
    args = ap.parse_args()

    import tensorflow as tf

    train = tf.keras.utils.image_dataset_from_directory(
        args.data, validation_split=0.2, subset="training", seed=42,
        image_size=IMG_SIZE, batch_size=args.batch)
    val = tf.keras.utils.image_dataset_from_directory(
        args.data, validation_split=0.2, subset="validation", seed=42,
        image_size=IMG_SIZE, batch_size=args.batch)
    class_names = train.class_names
    Path(args.labels).write_text("\n".join(class_names), encoding="utf-8")

    train = train.prefetch(tf.data.AUTOTUNE)
    val = val.prefetch(tf.data.AUTOTUNE)

    model = build_model(len(class_names))
    model.compile(optimizer="adam", loss="sparse_categorical_crossentropy",
                  metrics=["accuracy"])
    model.fit(train, validation_data=val, epochs=args.epochs)
    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    model.save(args.out)
    print(f"Saved {args.task} model -> {args.out} ({len(class_names)} classes); "
          f"labels -> {args.labels}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())