#
# @File: classify_views.py
# @Author: FlexDealer
# Copyright (c) 2026 FlexDealer
#
# Real-time-feed exterior-view classifier.
#
# Modes:
#   CLI:    python classify_views.py --urls-file paths.txt   (or a JSON array on stdin)
#   Daemon: python classify_views.py --serve --port 8788     (models load once, kept warm)
#           POST {"paths":[...]} to /classify ; GET /health -> {"status":"ready"}
#
# Pipeline: CLIP car/non-car (non-car discarded) -> CLIP interior/exterior (interior kept,
# no view) -> angle model assigns an exterior view -> classifier breaks ties within a view.
# CLIP needs transformers/torch; the angle and classifier models need TensorFlow.

import os
import sys
import json
import argparse
from collections import defaultdict

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODELS_DIR = os.path.join(BASE_DIR, "..", "models")
if BASE_DIR not in sys.path:            # so training_tools.py (angle custom objects) is importable
    sys.path.insert(0, BASE_DIR)

# Model paths (override via env for local testing).
ANGLE_PATH      = os.environ.get("VD_ANGLE",      os.path.join(MODELS_DIR, "efficientnetb0_approach2.h5"))
CLASSIFIER_PATH = os.environ.get("VD_CLASSIFIER", os.path.join(MODELS_DIR, "liftkit_custom_trained_model.keras"))
CLIP_DIR        = os.environ.get("VD_CLIP",       os.path.join(MODELS_DIR, "clip"))
CLASS_NAMES_FILE = os.path.join(BASE_DIR, "class_names.txt")

IMG_SIZE = (224, 224)

# CLIP gating: an image is CAR when car_mass >= CLIP_CAR_MIN, EXTERIOR when ext_mass >= CLIP_EXT_MIN.
# Thresholds sit inside measured gaps (cars car_mass >= 0.60 vs non-car <= 0.32; exteriors
# ext_mass >= 0.94 vs interiors <= 0.06).
CLIP_CAR_MIN = 0.45
CLIP_EXT_MIN = 0.50
CAR_PROMPTS = [
    "a photo of a car", "a car parked outdoors", "the exterior of a car", "a car in a parking lot",
    "a car on a road", "a vehicle photographed outdoors", "an automobile", "a parked vehicle",
    "a close-up of a car part", "the interior of a car", "a car wheel or tire",
    "a car engine bay", "a car dashboard", "a car seat",
]
NON_CAR_PROMPTS = [
    "a company logo or brand mark", "a website or advertising banner",
    "a document, chart, or screenshot", "text on a plain background",
    "a solid color or gradient background", "a product icon",
]
EXTERIOR_PROMPTS = [
    "the exterior of a car", "a car photographed from outside", "the front of a car",
    "the side of a car", "the rear of a car", "a car parked outdoors", "a whole car seen from outside",
]
INTERIOR_PROMPTS = [
    "the interior of a car", "the inside of a car cabin", "a car dashboard", "car seats",
    "a car steering wheel", "the cargo area inside a car", "a close-up inside a car",
]

# Canonical angle (deg) per exterior view; each predicted angle maps to the nearest one.
VIEW_CENTERS = {
    "eff": 0, "efd": 25, "esd": 75, "erd": 155,
    "erf": 182, "erp": 200, "esp": 270, "efp": 335,
}


def eprint(*a):
    print(*a, file=sys.stderr)


def circular_nearest_view(angle_deg):
    best_view, best_dist = None, 1e9
    for view, center in VIEW_CENTERS.items():
        d = abs((angle_deg - center + 180) % 360 - 180)
        if d < best_dist:
            best_dist, best_view = d, view
    return best_view


def angle_closeness(angle_deg, view):
    """Proxy confidence (0..1) for a view with a single candidate."""
    dist = abs((angle_deg - VIEW_CENTERS[view] + 180) % 360 - 180)
    return round(max(0.0, 1.0 - dist / 45.0), 3)


class ViewClassifier:
    """Loads CLIP, the angle model and the view classifier lazily, once per process."""

    def __init__(self):
        self._angle_model = None
        self._class_model = None
        self._clip_model = None
        self._clip_proc = None
        if os.path.exists(CLASS_NAMES_FILE):
            with open(CLASS_NAMES_FILE) as f:
                self._class_names = [ln.strip() for ln in f if ln.strip()]
        else:
            self._class_names = ["efd", "eff", "efp", "erd", "erf", "erp", "esd", "esp"]

    # ---------- Keras loading (works on TF 2.16 tf_keras and TF 2.12 bundled keras) ----------
    def _keras_load_fns(self):
        import importlib
        os.environ.setdefault("TF_USE_LEGACY_KERAS", "1")
        os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
        fns = []
        for path in ("tf_keras.models", "tensorflow.keras.models", "keras.models"):
            try:
                fns.append(importlib.import_module(path).load_model)
            except Exception:
                pass
        return fns

    def _custom_objects(self):
        custom = {}
        try:
            import tensorflow_hub as hub
            custom["KerasLayer"] = hub.KerasLayer
        except Exception:
            pass
        try:
            import training_tools as tt
            for n in ("tf_mean_absolute_angle_error_sin_cos_output", "tf_rmse_angle_sin_cos_output",
                      "tf_r2_angle_score_sin_cos_output", "tf_median_absolute_angle_error_sin_cos_output",
                      "tf_acc_pi_6_sin_cos_output", "horizontal_flip_pose_sin_cos_output"):
                if hasattr(tt, n):
                    custom[n] = getattr(tt, n)
        except Exception:
            pass
        return custom

    def _load_keras(self, path):
        custom = self._custom_objects()
        last = None
        for load in self._keras_load_fns():
            for kwargs in ({"custom_objects": custom, "compile": False}, {"compile": False}):
                try:
                    return load(path, **kwargs)
                except Exception as e:
                    last = e
        raise RuntimeError("could not load keras model %s: %s" % (path, last))

    def _angle(self):
        if self._angle_model is None:
            self._angle_model = self._load_keras(ANGLE_PATH)
        return self._angle_model

    def _classifier(self):
        if self._class_model is None:
            self._class_model = self._load_keras(CLASSIFIER_PATH)
        return self._class_model

    def _clip(self):
        if self._clip_model is None:
            from transformers import CLIPProcessor, CLIPModel
            self._clip_proc = CLIPProcessor.from_pretrained(CLIP_DIR)
            self._clip_model = CLIPModel.from_pretrained(CLIP_DIR)
            self._clip_model.eval()
        return self._clip_model, self._clip_proc

    # ---------- Stages 1-2: CLIP gating ----------
    def _clip_group_mass(self, rgb_img, group_a, group_b):
        """Softmax mass on group_a within (group_a + group_b) prompts (0..1)."""
        import torch
        model, proc = self._clip()
        with torch.no_grad():
            inp = proc(text=group_a + group_b, images=rgb_img, return_tensors="pt", padding=True)
            probs = model(**inp).logits_per_image[0].softmax(-1)
        return float(probs[:len(group_a)].sum())

    def _stage_gate(self, pil_img):
        """Return "non-car", "interior" or "exterior"."""
        rgb = pil_img.convert("RGB")
        if self._clip_group_mass(rgb, CAR_PROMPTS, NON_CAR_PROMPTS) < CLIP_CAR_MIN:
            return "non-car"
        if self._clip_group_mass(rgb, EXTERIOR_PROMPTS, INTERIOR_PROMPTS) >= CLIP_EXT_MIN:
            return "exterior"
        return "interior"

    @staticmethod
    def _img_array(pil_img):
        import numpy as np
        return np.array(pil_img.convert("RGB").resize(IMG_SIZE), dtype="float32")

    # ---------- Stage 3: angle -> view ----------
    def _predict_angle(self, pil_img):
        import numpy as np
        import tensorflow as tf
        out = self._angle().predict(tf.expand_dims(tf.convert_to_tensor(self._img_array(pil_img)), 0), verbose=0)
        angle = float(np.arctan2(out[0][0], out[0][1]) / np.pi * 180)   # sin, cos -> degrees
        angle = (angle + 180) % 360 - 180
        return (360 + angle) if angle < 0 else angle

    # ---------- Stage 4: classifier score per view (breaks ties within a bucket) ----------
    def _classifier_scores(self, pil_img):
        import numpy as np
        from tensorflow.keras.applications.resnet50 import preprocess_input
        pred = self._classifier().predict(preprocess_input(np.expand_dims(self._img_array(pil_img), 0)), verbose=0)[0]
        return {self._class_names[i]: float(pred[i]) for i in range(len(self._class_names))}

    @staticmethod
    def _load_image(fp):
        """Open an image and read it into memory so the file descriptor is closed on return."""
        from PIL import Image
        with Image.open(fp) as img:
            img.load()
            return img

    def _resolve(self, source):
        """Return (pil_image, local_path, is_temp); downloads URLs to a temp file."""
        from io import BytesIO
        if source.startswith("http://") or source.startswith("https://"):
            import requests, tempfile
            r = requests.get(source, timeout=60)
            r.raise_for_status()
            fd, tmp = tempfile.mkstemp(suffix=".jpg")
            try:
                with os.fdopen(fd, "wb") as fh:
                    fh.write(r.content)
                return self._load_image(BytesIO(r.content)), tmp, True
            except Exception:
                try:
                    os.remove(tmp)
                except Exception:
                    pass
                raise
        p = os.path.expanduser(source)
        return self._load_image(p), p, False

    def classify_one(self, source):
        try:
            pil, path, is_temp = self._resolve(source)
        except Exception as e:
            return {"view": None, "confidence": 0.0, "angle": None, "note": "error: %s" % e,
                    "_path": None, "_tmp": False}
        try:
            stage = self._stage_gate(pil)
            if stage != "exterior":     # non-car -> caller discards; interior -> kept, no view
                return {"view": None, "confidence": 0.0, "angle": None, "note": stage,
                        "_path": path, "_tmp": is_temp}
            angle = self._predict_angle(pil)
            view = circular_nearest_view(angle)
            return {"view": view, "confidence": angle_closeness(angle, view), "angle": round(angle, 1),
                    "note": "exterior", "_path": path, "_tmp": is_temp}
        except Exception as e:
            return {"view": None, "confidence": 0.0, "angle": None, "note": "error: %s" % e,
                    "_path": path, "_tmp": is_temp}

    def classify_batch(self, sources):
        raw = [(s, self.classify_one(s)) for s in sources]

        # Group exteriors by view; one candidate keeps its angle confidence, several are ranked
        # by the classifier so only the best per view keeps a view code.
        buckets = defaultdict(list)
        for s, res in raw:
            if res["view"]:
                buckets[res["view"]].append(res)

        winners = set()
        for view, cands in buckets.items():
            if len(cands) == 1:
                winners.add(id(cands[0]))
                continue
            best, best_score = None, -1.0
            for res in cands:
                try:
                    score = self._classifier_scores(self._load_image(res["_path"]))[view]
                except Exception:
                    score = res["confidence"]
                res["confidence"] = round(float(score), 3)
                if score > best_score:
                    best_score, best = score, res
            winners.add(id(best))

        for _, res in raw:
            if res["view"] and id(res) not in winners:
                res["note"] = "duplicate-view-demoted"
                res["view"] = None
            if res.get("_tmp") and res.get("_path") and os.path.exists(res["_path"]):
                try:
                    os.remove(res["_path"])
                except Exception:
                    pass

        return [{"url": s, "view": res["view"], "confidence": res["confidence"],
                 "angle": res["angle"], "note": res["note"]} for s, res in raw]


def serve(host, port, idle):
    import threading, time
    from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

    clf = ViewClassifier()
    try:
        clf._clip()                 # pre-warm so the first request isn't slow
    except Exception as e:
        eprint("CLIP pre-warm failed (will retry lazily): %s" % e)
    lock = threading.Lock()
    last = [time.time()]

    class Handler(BaseHTTPRequestHandler):
        def log_message(self, *a):
            pass

        def _reply(self, obj, code=200):
            body = json.dumps(obj).encode()
            self.send_response(code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self):
            if self.path.startswith("/health"):
                self._reply({"status": "ready"})
            else:
                self._reply({"status": "error", "message": "not found"}, 404)

        def do_POST(self):
            try:
                n = int(self.headers.get("Content-Length", 0))
                paths = (json.loads(self.rfile.read(n) or b"{}").get("paths")) or []
            except Exception as e:
                self._reply({"status": "error", "message": "bad request: %s" % e}, 400)
                return
            with lock:                  # inference is serialized (TF/torch are not thread-safe)
                results = clf.classify_batch(paths)
            last[0] = time.time()
            self._reply({"status": "success", "results": results})

    srv = ThreadingHTTPServer((host, port), Handler)

    def watchdog():
        while idle:
            time.sleep(30)
            if time.time() - last[0] > idle:
                srv.shutdown()
                return
    if idle:
        threading.Thread(target=watchdog, daemon=True).start()

    eprint("classify_views daemon ready on %s:%d (idle=%ss)" % (host, port, idle))
    srv.serve_forever()


def read_sources(args):
    if args.urls_file:
        with open(args.urls_file) as f:
            return [ln.strip() for ln in f if ln.strip()]
    data = sys.stdin.read().strip()
    return json.loads(data) if data else []


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--serve", action="store_true", help="run as a persistent HTTP daemon")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8788)
    ap.add_argument("--idle", type=int, default=600, help="daemon idle-shutdown seconds (0=never)")
    ap.add_argument("--urls-file", help="one image URL/path per line; else JSON array on stdin")
    args = ap.parse_args()

    if args.serve:
        serve(args.host, args.port, args.idle)
        return

    try:
        sources = read_sources(args)
    except Exception as e:
        print(json.dumps({"status": "error", "message": "bad input: %s" % e, "results": []}))
        return
    if not sources:
        print(json.dumps({"status": "success", "results": []}))
        return
    try:
        results = ViewClassifier().classify_batch(sources)
        print(json.dumps({"status": "success", "results": results}))
    except Exception as e:
        print(json.dumps({"status": "error", "message": str(e), "results": []}))


if __name__ == "__main__":
    main()