"""
Utilities for training a FOMO-style (Faster Objects, More Objects) object
detection model on a custom dataset exported from Roboflow in COCO JSON
format, and quantizing it to int8 TFLite for deployment on a
Seeed XIAO ESP32S3 Sense.

FOMO reframes detection as a per-grid-cell classification problem: instead
of predicting bounding boxes, the model predicts an object class (or
"background") for each cell of a coarse output grid (stride 8 relative to
the input). This is dramatically cheaper than YOLO/SSD-style detection and
is the standard approach used for MCU-class hardware like the ESP32S3.
"""

import os
import json
import zipfile
import glob

import numpy as np
import tensorflow as tf
from PIL import Image, ImageDraw


# --------------------------------------------------------------------------
# Dataset extraction / parsing
# --------------------------------------------------------------------------

def extract_zip(zip_path, dest_dir):
    os.makedirs(dest_dir, exist_ok=True)
    with zipfile.ZipFile(zip_path, "r") as zf:
        zf.extractall(dest_dir)
    return dest_dir


def find_coco_json(root_dir):
    candidates = glob.glob(os.path.join(root_dir, "**", "*.json"), recursive=True)
    for path in candidates:
        try:
            with open(path, "r") as f:
                data = json.load(f)
            if "images" in data and "annotations" in data and "categories" in data:
                return path, data
        except (json.JSONDecodeError, UnicodeDecodeError):
            continue
    raise FileNotFoundError(
        f"No COCO-format _annotations.coco.json found under {root_dir}. "
        "Zip the split folder exactly as Roboflow exported it "
        "(images + _annotations.coco.json together, no extra nesting)."
    )


def load_coco_split(zip_path, work_dir, split_name):
    """Extract a Roboflow COCO zip and return (samples, category_id_to_name)."""
    dest = os.path.join(work_dir, split_name)
    extract_zip(zip_path, dest)
    json_path, coco = find_coco_json(dest)
    images_dir = os.path.dirname(json_path)

    id_to_file = {img["id"]: img["file_name"] for img in coco["images"]}
    categories = {c["id"]: c["name"] for c in coco["categories"]}

    anns_by_image = {}
    for ann in coco["annotations"]:
        anns_by_image.setdefault(ann["image_id"], []).append(ann)

    samples = []
    for img_id, file_name in id_to_file.items():
        img_path = os.path.join(images_dir, file_name)
        if not os.path.exists(img_path):
            alt = os.path.join(images_dir, "images", file_name)
            img_path = alt if os.path.exists(alt) else img_path
        samples.append((img_path, anns_by_image.get(img_id, [])))

    return samples, categories


def load_coco_split_from_files(file_objs, split_name):
    """Build a split from a loose list of uploaded files (a folder upload
    containing the images plus one _annotations.coco.json), rather than a
    zip. Matches images to annotation entries by basename, since browser
    directory uploads don't always preserve relative folder paths.
    """
    if not file_objs:
        raise FileNotFoundError(
            f"No files were uploaded for the '{split_name}' split. "
            "Select the folder containing your images and _annotations.coco.json."
        )

    json_data = None
    image_paths = {}  # basename -> path on disk
    for f in file_objs:
        path = f.name if hasattr(f, "name") else f
        base = os.path.basename(path)
        if base.lower().endswith(".json"):
            try:
                with open(path, "r") as fh:
                    candidate = json.load(fh)
                if all(k in candidate for k in ("images", "annotations", "categories")):
                    json_data = candidate
            except (json.JSONDecodeError, UnicodeDecodeError):
                continue
        else:
            image_paths[base] = path

    if json_data is None:
        raise FileNotFoundError(
            f"No COCO-format _annotations.coco.json found among the uploaded "
            f"'{split_name}' files. Make sure it's included in the folder you selected."
        )

    id_to_file = {img["id"]: img["file_name"] for img in json_data["images"]}
    categories = {c["id"]: c["name"] for c in json_data["categories"]}

    anns_by_image = {}
    for ann in json_data["annotations"]:
        anns_by_image.setdefault(ann["image_id"], []).append(ann)

    def normalize(name):
        stem = os.path.splitext(name)[0]
        return "".join(ch for ch in stem.lower() if ch.isalnum())

    normalized_lookup = {normalize(base): path for base, path in image_paths.items()}

    samples = []
    missing = 0
    for img_id, file_name in id_to_file.items():
        base = os.path.basename(file_name)
        img_path = image_paths.get(base)
        if img_path is None:
            # Fall back to a normalized match in case dots/punctuation in the
            # filename got changed somewhere along the upload path.
            img_path = normalized_lookup.get(normalize(base))
        if img_path is None:
            missing += 1
            continue
        samples.append((img_path, anns_by_image.get(img_id, [])))

    if not samples:
        raise FileNotFoundError(
            f"The '{split_name}' JSON references images, but none of the uploaded "
            "files matched them by name. Double check every image is included."
        )

    return samples, categories, missing


def build_label_map(train_categories, test_categories):
    names = sorted(set(train_categories.values()) | set(test_categories.values()))
    names = [n for n in names if n.strip().lower() not in ("background", "objects", "")]
    if not names:
        return {}
    return {name: i + 1 for i, name in enumerate(names)}  # 0 is reserved for background


# --------------------------------------------------------------------------
# tf.data pipeline
# --------------------------------------------------------------------------

def make_dataset(samples, label_map, input_size, grid_size, batch_size, augment, cat_id_to_name):
    stride = input_size / grid_size

    def gen():
        for img_path, anns in samples:
            try:
                img = Image.open(img_path).convert("RGB")
            except Exception:
                continue
            orig_w, orig_h = img.size
            img_resized = img.resize((input_size, input_size))
            arr = np.asarray(img_resized, dtype=np.float32) / 255.0

            label_grid = np.zeros((grid_size, grid_size), dtype=np.int32)
            for ann in anns:
                name = cat_id_to_name.get(ann["category_id"])
                cls_idx = label_map.get(name)
                if cls_idx is None:
                    continue
                x, y, w, h = ann["bbox"]  # COCO: top-left x, y, width, height
                cx = (x + w / 2) / orig_w * input_size
                cy = (y + h / 2) / orig_h * input_size
                gx = min(int(cx // stride), grid_size - 1)
                gy = min(int(cy // stride), grid_size - 1)
                label_grid[gy, gx] = cls_idx

            yield arr, label_grid

    ds = tf.data.Dataset.from_generator(
        gen,
        output_signature=(
            tf.TensorSpec(shape=(input_size, input_size, 3), dtype=tf.float32),
            tf.TensorSpec(shape=(grid_size, grid_size), dtype=tf.int32),
        ),
    )

    if augment:
        def aug(img, label):
            if tf.random.uniform(()) > 0.5:
                img = tf.image.flip_left_right(img)
                label = tf.reverse(label, axis=[1])
            img = tf.image.random_brightness(img, 0.2)
            img = tf.clip_by_value(img, 0.0, 1.0)
            return img, label
        ds = ds.map(aug, num_parallel_calls=tf.data.AUTOTUNE)

    ds = ds.shuffle(256).repeat().batch(batch_size).prefetch(tf.data.AUTOTUNE)
    return ds


# --------------------------------------------------------------------------
# Model
# --------------------------------------------------------------------------

def build_fomo_model(input_size, num_classes, alpha=0.35):
    """MobileNetV2 backbone cut at 1/8 resolution + a small conv head.
    This mirrors the architecture commonly called 'FOMO'."""
    inputs = tf.keras.Input(shape=(input_size, input_size, 3))
    base = tf.keras.applications.MobileNetV2(
        input_tensor=inputs, alpha=alpha, include_top=False, weights="imagenet"
    )
    target_size = input_size // 8
    cut_layer = None
    for layer in base.layers:
        try:
            shape = layer.output.shape
        except AttributeError:
            continue
        if shape is None or len(shape) != 4:
            continue
        if shape[1] == target_size:
            cut_layer = layer
    if cut_layer is None:
        raise ValueError(f"Could not find a layer with spatial size {target_size} in MobileNetV2.")

    x = cut_layer.output
    x = tf.keras.layers.Conv2D(32, 1, activation="relu", name="fomo_head_conv")(x)
    outputs = tf.keras.layers.Conv2D(
        num_classes + 1, 1, activation="softmax", name="fomo_output"
    )(x)
    return tf.keras.Model(inputs, outputs, name="fomo")


def get_optimizer(name, learning_rate):
    name = (name or "adamw").lower()
    if name == "adamw":
        return tf.keras.optimizers.AdamW(learning_rate=learning_rate, weight_decay=1e-4)
    if name == "sgd":
        return tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=0.9)
    return tf.keras.optimizers.Adam(learning_rate=learning_rate)


def weighted_sparse_ce(background_weight):
    def loss_fn(y_true, y_pred):
        y_true = tf.cast(y_true, tf.int32)
        per_cell = tf.keras.losses.sparse_categorical_crossentropy(y_true, y_pred)
        weights = tf.where(tf.equal(y_true, 0), background_weight, 1.0)
        return tf.reduce_mean(per_cell * weights)
    return loss_fn


# --------------------------------------------------------------------------
# Training
# --------------------------------------------------------------------------

class LogCallback(tf.keras.callbacks.Callback):
    """Appends per-epoch log lines to a list so the caller can stream them."""

    def __init__(self, log_lines):
        super().__init__()
        self.log_lines = log_lines

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        line = f"Epoch {epoch + 1}: " + ", ".join(f"{k}={v:.4f}" for k, v in logs.items())
        self.log_lines.append(line)


def train_model(model, train_ds, val_ds, epochs, learning_rate, background_weight, log_lines,
                 optimizer_name="adamw", steps_per_epoch=None, validation_steps=None):
    model.compile(
        optimizer=get_optimizer(optimizer_name, learning_rate),
        loss=weighted_sparse_ce(background_weight),
        metrics=["accuracy"],
    )
    callback = LogCallback(log_lines)
    history = model.fit(
        train_ds,
        validation_data=val_ds,
        epochs=epochs,
        steps_per_epoch=steps_per_epoch,
        validation_steps=validation_steps,
        callbacks=[callback],
        verbose=0,
    )
    return history


# --------------------------------------------------------------------------
# Quantization
# --------------------------------------------------------------------------

def quantize_model(model, representative_samples, input_size):
    def rep_dataset():
        for img_path, _ in representative_samples[:100]:
            try:
                img = Image.open(img_path).convert("RGB").resize((input_size, input_size))
            except Exception:
                continue
            arr = np.asarray(img, dtype=np.float32) / 255.0
            yield [arr[np.newaxis, ...]]

    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.representative_dataset = rep_dataset
    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
    converter.inference_input_type = tf.int8
    converter.inference_output_type = tf.int8
    return converter.convert()


def tflite_to_c_header(tflite_bytes, var_name="crow_fomo_model"):
    lines = [
        "// Auto-generated TFLite model as a C array for ESP32 deployment",
        f"#ifndef {var_name.upper()}_H",
        f"#define {var_name.upper()}_H",
        "",
        f"alignas(8) const unsigned char {var_name}[] = {{",
    ]
    hex_bytes = [f"0x{b:02x}" for b in tflite_bytes]
    for i in range(0, len(hex_bytes), 12):
        lines.append("  " + ", ".join(hex_bytes[i:i + 12]) + ",")
    lines.append("};")
    lines.append(f"const unsigned int {var_name}_len = {len(tflite_bytes)};")
    lines.append("")
    lines.append("#endif")
    return "\n".join(lines)


# --------------------------------------------------------------------------
# Inference (for the "Test Model" tab)
# --------------------------------------------------------------------------

def run_tflite_inference(tflite_path, pil_image, label_names, threshold=0.5):
    interpreter = tf.lite.Interpreter(model_path=tflite_path)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()[0]
    output_details = interpreter.get_output_details()[0]

    input_size = input_details["shape"][1]
    grid_size = output_details["shape"][1]

    orig_w, orig_h = pil_image.size
    img_resized = pil_image.convert("RGB").resize((input_size, input_size))
    arr = np.asarray(img_resized, dtype=np.float32) / 255.0

    in_scale, in_zero = input_details["quantization"]
    if in_scale > 0:
        arr_q = (arr / in_scale + in_zero).astype(np.int8)
    else:
        arr_q = arr.astype(np.int8)
    interpreter.set_tensor(input_details["index"], arr_q[np.newaxis, ...])
    interpreter.invoke()
    out = interpreter.get_tensor(output_details["index"])[0]

    out_scale, out_zero = output_details["quantization"]
    if out_scale > 0:
        out = (out.astype(np.float32) - out_zero) * out_scale

    stride_x = orig_w / grid_size
    stride_y = orig_h / grid_size

    draw_img = pil_image.convert("RGB").copy()
    draw = ImageDraw.Draw(draw_img)
    detections = []

    for gy in range(grid_size):
        for gx in range(grid_size):
            cell = out[gy, gx]
            cls_idx = int(np.argmax(cell))
            conf = float(cell[cls_idx])
            if cls_idx == 0 or conf < threshold:
                continue
            cx = (gx + 0.5) * stride_x
            cy = (gy + 0.5) * stride_y
            r = min(stride_x, stride_y) * 0.6
            draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=(255, 0, 0), width=3)
            label_name = label_names[cls_idx - 1] if cls_idx - 1 < len(label_names) else str(cls_idx)
            draw.text((cx + r, cy - r), f"{label_name} {conf:.2f}", fill=(255, 0, 0))
            detections.append({"class": label_name, "confidence": conf, "x": cx, "y": cy})

    return draw_img, detections