#!/usr/bin/env python3
"""
update_site_metadata_from_vector_class_evidence.py

Vector-class aware orientation updater.

This prevents a hall/building axis from being accidentally treated as a machine
physics vector.

Promotion logic:
LOCATION_ONLY_NOT_DIAGNOSTIC
  -> HALL_AXIS_CANDIDATE
  -> MACHINE_AXIS_CANDIDATE
  -> TRANSFER_PENDING
  -> DIAGNOSTIC_READY

Building/hall axes are stored in hall_axis_azimuth_deg only.
Machine/physics axes are stored in machine_axis_azimuth_deg only when
axis_is_physics_vector_true_false is true.
"""
from pathlib import Path
import pandas as pd
import numpy as np
import math

ROOT = Path(__file__).resolve().parents[1]
site_path = ROOT / "data" / "site_metadata" / "tokamak_site_metadata_v1_vector_class_schema.csv"
evidence_path = ROOT / "data" / "orientation_evidence" / "orientation_evidence_log_v2_FILLED.csv"
out_path = ROOT / "data" / "site_metadata" / "tokamak_site_metadata_v2_vector_class_candidates.csv"

PHYSICS_VECTOR_CLASSES = {
    "device_vertical_axis",
    "neutral_beam_axis",
    "diagnostic_port_axis",
    "toroidal_field_reference_direction",
    "machine_layout_reference_axis",
}

HALL_ONLY_CLASSES = {
    "torus_hall_structural_axis",
}

def bearing(lat1, lon1, lat2, lon2):
    phi1 = math.radians(float(lat1))
    phi2 = math.radians(float(lat2))
    dlon = math.radians(float(lon2) - float(lon1))
    x = math.sin(dlon) * math.cos(phi2)
    y = math.cos(phi1) * math.sin(phi2) - math.sin(phi1) * math.cos(phi2) * math.cos(dlon)
    return (math.degrees(math.atan2(x, y)) + 360.0) % 360.0

def parse_bool(v):
    if isinstance(v, bool):
        return v
    s = str(v).strip().lower()
    return s in ["true", "1", "yes", "y"]

def get_axis_azimuth(row):
    az = pd.to_numeric(row.get("manual_axis_azimuth_deg"), errors="coerce")
    if not pd.isna(az):
        return float(az) % 360.0
    lat1 = pd.to_numeric(row.get("axis_point_A_lat_deg"), errors="coerce")
    lon1 = pd.to_numeric(row.get("axis_point_A_lon_deg"), errors="coerce")
    lat2 = pd.to_numeric(row.get("axis_point_B_lat_deg"), errors="coerce")
    lon2 = pd.to_numeric(row.get("axis_point_B_lon_deg"), errors="coerce")
    if not any(pd.isna(v) for v in [lat1, lon1, lat2, lon2]):
        return bearing(lat1, lon1, lat2, lon2)
    return np.nan

site = pd.read_csv(site_path)
if not evidence_path.exists():
    raise SystemExit(f"Missing filled evidence file: {evidence_path}")

ev = pd.read_csv(evidence_path)

for _, row in ev.iterrows():
    tok = str(row.get("TOK")).strip()
    vector_class = str(row.get("orientation_vector_class", "unknown")).strip()
    is_physics = parse_bool(row.get("axis_is_physics_vector_true_false", False))
    conf = pd.to_numeric(row.get("axis_source_confidence_0_1"), errors="coerce")
    if pd.isna(conf) or conf <= 0:
        continue
    az = get_axis_azimuth(row)
    if pd.isna(az):
        continue

    mask = site["TOK"].astype(str).eq(tok)
    if not mask.any():
        continue

    if vector_class in HALL_ONLY_CLASSES or not is_physics:
        site.loc[mask, "hall_axis_azimuth_deg"] = az
        site.loc[mask, "axis_source_confidence_0_1"] = max(float(conf), float(site.loc[mask, "axis_source_confidence_0_1"].fillna(0).iloc[0]))
        site.loc[mask, "orientation_vector_class"] = vector_class
        site.loc[mask, "axis_is_physics_vector_true_false"] = False
        site.loc[mask, "axis_evidence_id"] = str(row.get("evidence_id", ""))
        site.loc[mask, "axis_promotion_status"] = "HALL_AXIS_CANDIDATE"
        site.loc[mask, "diagnostic_status"] = "HALL_AXIS_ONLY_MACHINE_AXIS_UNVERIFIED"
        site.loc[mask, "needs_next"] = "machine/physics vector still required"
    elif vector_class in PHYSICS_VECTOR_CLASSES and is_physics:
        site.loc[mask, "machine_axis_azimuth_deg"] = az
        site.loc[mask, "device_axis_azimuth_deg"] = az
        site.loc[mask, "device_axis_tilt_deg"] = pd.to_numeric(row.get("manual_axis_tilt_deg"), errors="coerce")
        site.loc[mask, "orientation_confidence_0_1"] = float(conf)
        site.loc[mask, "axis_source_confidence_0_1"] = float(conf)
        site.loc[mask, "orientation_vector_class"] = vector_class
        site.loc[mask, "axis_is_physics_vector_true_false"] = True
        site.loc[mask, "axis_evidence_id"] = str(row.get("evidence_id", ""))
        site.loc[mask, "axis_promotion_status"] = "MACHINE_AXIS_CANDIDATE"
        site.loc[mask, "diagnostic_status"] = "ORIENTATION_CANDIDATE_TRANSFER_PENDING"
        site.loc[mask, "needs_next"] = "gamma_lab, gamma_wall, gamma_b and transfer confidence required"

for col in ["location_confidence_0_1", "orientation_confidence_0_1", "transfer_confidence_0_1"]:
    site[col] = pd.to_numeric(site[col], errors="coerce").fillna(0)

site["metadata_confidence_0_1"] = site["location_confidence_0_1"] * site["orientation_confidence_0_1"] * site["transfer_confidence_0_1"]
site.to_csv(out_path, index=False)
print(f"Wrote {out_path}")
print(site[["TOK", "hall_axis_azimuth_deg", "machine_axis_azimuth_deg", "orientation_vector_class", "axis_is_physics_vector_true_false", "axis_promotion_status", "diagnostic_status", "metadata_confidence_0_1"]].to_string(index=False))
