#!/usr/bin/env python3
"""
Radiacode 103 — Identification automatique d'isotopes
Cycle de 24h avec détection branché/débranché
Fonctionne en Docker sur machine GPU (dev) ou RPi 4 (production)
"""

import numpy as np
import torch
import time
import json
import logging
import os
import sys
from datetime import datetime
from pathlib import Path

# Configuration via variables d'environnement
MODEL_PATH = os.environ.get("MODEL_PATH", "/models/vega_best.pt")
ISOTOPE_INDEX_PATH = os.environ.get("ISOTOPE_INDEX_PATH", "/models/vega_isotope_index.txt")
BACKGROUND_PATH = os.environ.get("BACKGROUND_PATH", "/data/background_24h.npy")
LOG_DIR = Path(os.environ.get("LOG_DIR", "/logs"))
LOG_DIR.mkdir(parents=True, exist_ok=True)
THRESHOLD = float(os.environ.get("THRESHOLD", "0.5"))
SAMPLE_INTERVAL = int(os.environ.get("SAMPLE_INTERVAL", "60"))
REPORT_HOUR = int(os.environ.get("REPORT_HOUR", "0"))
MIN_LIVE_TIME = int(os.environ.get("MIN_LIVE_TIME", "3600"))
STATE_PATH = os.environ.get("STATE_PATH", "/data/monitor_state.json")
CPS_LOG_PATH = os.environ.get("CPS_LOG_PATH", "/data/cps_log.jsonl")
ENERGY_OFFSET = float(os.environ.get("ENERGY_CALIBRATION_OFFSET", "0.33"))
ENERGY_SLOPE = float(os.environ.get("ENERGY_CALIBRATION_SLOPE", "2.97"))

# Logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(LOG_DIR / "radiacode.log"),
    ],
)
log = logging.getLogger(__name__)


class RadiacodeMonitor:
    def __init__(self):
        # Charger le modèle PyTorch
        device_str = os.environ.get("VEGA_DEVICE", "cpu")
        self.device = torch.device(device_str)

        log.info(f"Chargement du modèle depuis {MODEL_PATH} sur {self.device}...")
        checkpoint = torch.load(MODEL_PATH, map_location=self.device, weights_only=False)

        # Importer VegaModel (depuis le volume monté)
        vega_ml_path = os.environ.get("VEGA_ML_PATH", "/models/vega_ml")
        if vega_ml_path not in sys.path:
            sys.path.insert(0, vega_ml_path)

        from training.vega.model import VegaModel, VegaConfig
        from training.vega.isotope_index import IsotopeIndex

        self.model_config = VegaConfig(**checkpoint["model_config"])
        self.model = VegaModel(self.model_config)
        self.model.load_state_dict(checkpoint["model_state_dict"])
        self.model.eval()
        log.info(
            f"Modèle chargé : {self.model_config.num_isotopes} isotopes, "
            f"{self.model.count_parameters():,} paramètres"
        )

        # Charger l'index des isotopes
        self.isotope_index = IsotopeIndex.load(Path(ISOTOPE_INDEX_PATH))

        # Charger le bruit de fond de référence
        self.bg_counts = None
        self.bg_live_time = None
        bg_path = Path(BACKGROUND_PATH)
        if bg_path.exists():
            bg_data = np.load(str(bg_path), allow_pickle=True).item()
            self.bg_counts = bg_data["counts"].astype(np.float64)
            self.bg_live_time = float(bg_data["duration"])
            log.info(
                f"Background chargé : {self.bg_live_time/3600:.1f}h, "
                f"{self.bg_counts.sum():.0f} coups"
            )
        else:
            log.warning(f"Pas de fichier background : {BACKGROUND_PATH}")

        # Compteurs cumulés
        self.cumulated_counts = np.zeros(1024, dtype=np.float64)
        self.cumulated_live_time = 0.0
        self.last_report_date = None
        self.connected = False

    def try_connect(self):
        """Tente de se connecter au Radiacode. Retourne le device ou None."""
        try:
            from radiacode import RadiaCode

            device = RadiaCode()
            log.info("Radiacode connecté")
            self.connected = True
            return device
        except Exception as e:
            log.debug(f"Détecteur non disponible : {e}")
            self.connected = False
            return None

    def sample_once(self):
        """Échantillonne une fois. Retourne True si succès."""
        device = None
        try:
            device = self.try_connect()
            if device is None:
                return False

            spectrum = device.spectrum()
            counts = np.array(spectrum.counts, dtype=np.float64)
            live_time = spectrum.duration.total_seconds()

            if live_time > 0 and counts.sum() > 0:
                self.cumulated_counts += counts
                self.cumulated_live_time += live_time
                device.spectrum_reset()
                log.info(
                    f"Échantillon : {counts.sum():.0f} coups en {live_time:.1f}s "
                    f"(cumul : {self.cumulated_live_time/3600:.1f}h)"
                )
                return True
            return False
        except Exception as e:
            log.warning(f"Erreur échantillonnage : {e}")
            return False
        finally:
            if device:
                try:
                    del device
                except Exception:
                    pass

    def save_state(self):
        """Ecrit l'etat actuel du moniteur dans un fichier JSON atomique."""
        energy_kev = [round(ENERGY_OFFSET + ENERGY_SLOPE * i, 2) for i in range(1024)]
        cps = float(self.cumulated_counts.sum() / self.cumulated_live_time) if self.cumulated_live_time > 0 else 0.0

        isotopes = []
        if self.cumulated_live_time > 0:
            rate = self.cumulated_counts / self.cumulated_live_time
            if self.bg_counts is not None and self.bg_live_time is not None:
                bg_rate = self.bg_counts / self.bg_live_time
                net_rate = np.clip(rate - bg_rate, 0, None)
            else:
                net_rate = rate
            isotopes = self.run_inference(net_rate)

        state = {
            "timestamp": datetime.now().isoformat(),
            "connected": self.connected,
            "cumulated_live_time_s": round(self.cumulated_live_time, 1),
            "cumulated_live_time_h": round(self.cumulated_live_time / 3600, 2),
            "total_counts": int(self.cumulated_counts.sum()),
            "cps": round(cps, 2),
            "background_subtracted": self.bg_counts is not None,
            "isotopes_detected": isotopes,
            "energy_kev": energy_kev,
            "counts": [round(float(c), 1) for c in self.cumulated_counts],
        }

        state_path = Path(STATE_PATH)
        state_path.parent.mkdir(parents=True, exist_ok=True)
        tmp_path = state_path.with_suffix(".tmp")
        with open(tmp_path, "w") as f:
            json.dump(state, f)
        os.replace(tmp_path, state_path)

    def log_cps(self, counts, live_time):
        """Ajoute un point CPS au journal horodaté."""
        cps = float(counts.sum() / live_time) if live_time > 0 else 0.0
        entry = {
            "ts": datetime.now().isoformat(),
            "cps": round(cps, 2),
            "live_time_s": round(live_time, 1),
            "total_counts": int(counts.sum()),
        }
        log_path = Path(CPS_LOG_PATH)
        log_path.parent.mkdir(parents=True, exist_ok=True)
        with open(log_path, "a") as f:
            f.write(json.dumps(entry) + "\n")

    def run_inference(self, spectrum_rate):
        """Exécute l'inférence PyTorch sur le spectre cumulé."""
        if spectrum_rate.max() > 0:
            normalized = spectrum_rate / spectrum_rate.max()
        else:
            return []

        tensor = torch.tensor(normalized, dtype=torch.float32).unsqueeze(0).to(self.device)

        with torch.no_grad():
            logits, activities = self.model(tensor)

        probs = torch.sigmoid(logits).cpu().numpy()[0]
        activities = activities.cpu().numpy()[0] * self.model_config.max_activity_bq

        results = []
        for i in range(len(probs)):
            if probs[i] >= THRESHOLD:
                results.append(
                    {
                        "isotope": self.isotope_index.index_to_name(i),
                        "probability": float(probs[i]),
                        "activity_bq": float(activities[i]),
                    }
                )

        return sorted(results, key=lambda x: -x["probability"])

    def generate_report(self):
        """Génère le rapport quotidien."""
        if self.cumulated_live_time < MIN_LIVE_TIME:
            log.warning(
                f"Pas assez de données ({self.cumulated_live_time/3600:.1f}h < "
                f"{MIN_LIVE_TIME/3600:.1f}h minimum). Pas de rapport."
            )
            return

        rate = self.cumulated_counts / self.cumulated_live_time

        if self.bg_counts is not None and self.bg_live_time is not None:
            bg_rate = self.bg_counts / self.bg_live_time
            net_rate = np.clip(rate - bg_rate, 0, None)
        else:
            net_rate = rate

        results = self.run_inference(net_rate)

        now = datetime.now()
        report = {
            "date": now.isoformat(),
            "live_time_hours": self.cumulated_live_time / 3600,
            "total_counts": int(self.cumulated_counts.sum()),
            "cps_mean": float(self.cumulated_counts.sum() / self.cumulated_live_time),
            "background_subtracted": self.bg_counts is not None,
            "isotopes_detected": results,
        }

        report_path = LOG_DIR / f"report_{now.strftime('%Y-%m-%d')}.json"
        with open(report_path, "w") as f:
            json.dump(report, f, indent=2, ensure_ascii=False)

        # Affichage
        print(f"\n{'='*50}")
        print(f"  RAPPORT — {now.strftime('%d/%m/%Y')}")
        print(f"{'='*50}")
        print(f"  Live time : {self.cumulated_live_time/3600:.1f}h")
        print(f"  Comptages : {self.cumulated_counts.sum():.0f}")
        print(f"  CPS moyen : {self.cumulated_counts.sum()/self.cumulated_live_time:.1f}")
        print(
            f"  Background : {'soustrait' if self.bg_counts is not None else 'non soustrait'}"
        )
        print()
        if results:
            for r in results:
                print(
                    f"  {r['isotope']:>10s} : {r['probability']*100:5.1f}% — {r['activity_bq']:.1f} Bq"
                )
        else:
            print("  (background uniquement)")
        print(f"{'='*50}\n")

        log.info(f"Rapport sauvegardé : {report_path}")

        # Reset pour le cycle suivant
        self.cumulated_counts = np.zeros(1024, dtype=np.float64)
        self.cumulated_live_time = 0.0

    def run(self):
        """Boucle principale."""
        log.info("=" * 50)
        log.info("Radiacode 103 — Moniteur d'isotopes")
        log.info("=" * 50)
        log.info(f"Modèle : {MODEL_PATH}")
        log.info(f"Device : {self.device}")
        log.info(f"Isotopes : {self.isotope_index.num_isotopes}")
        log.info(
            f"Background : {'chargé' if self.bg_counts is not None else 'non disponible'}"
        )
        log.info(f"Seuil : {THRESHOLD}")
        log.info(f"Intervalle : {SAMPLE_INTERVAL}s")

        while True:
            now = datetime.now()

            if self.last_report_date != now.date() and now.hour == REPORT_HOUR:
                self.generate_report()
                self.last_report_date = now.date()

            success = self.sample_once()
            if success:
                self.save_state()
                self.log_cps(self.cumulated_counts, self.cumulated_live_time)
            else:
                self.save_state()
            time.sleep(SAMPLE_INTERVAL)


if __name__ == "__main__":
    monitor = RadiacodeMonitor()
    monitor.run()