import gradio as gr
import json
import pandas as pd
import plotly.express as px
import tempfile
import numpy as np
import matplotlib.pyplot as plt


# -----------------------------
# Haversine distance (km)
# -----------------------------
def haversine(lat1, lon1, lat2, lon2):
    R = 6371.0
    phi1, phi2 = np.radians(lat1), np.radians(lat2)
    dphi = np.radians(lat2 - lat1)
    dlambda = np.radians(lon2 - lon1)

    a = np.sin(dphi / 2) ** 2 + np.cos(phi1) * np.cos(phi2) * np.sin(dlambda / 2) ** 2
    return 2 * R * np.arcsin(np.sqrt(a))


# -----------------------------
# Angle difference helper
# -----------------------------
def angle_diff(a, b):
    diff = abs(a - b)
    return min(diff, 360 - diff)


# -----------------------------
# Robust timestamp parser
# -----------------------------
def parse_timestamp(series):
    ts = pd.to_datetime(series, errors="coerce", utc=True)

    mask = ts.isna()
    if mask.any():
        s = series[mask].dropna()
        numeric = pd.to_numeric(s, errors="coerce")

        if not numeric.empty:
            if numeric.median() > 1e12:
                ts.loc[mask] = pd.to_datetime(numeric, unit="ms", errors="coerce", utc=True)
            else:
                ts.loc[mask] = pd.to_datetime(numeric, unit="s", errors="coerce", utc=True)

    return ts


# -----------------------------
# Load MMSIs
# -----------------------------
def load_geojson(geojson_file):
    if geojson_file is None:
        return gr.update(choices=[])

    with open(geojson_file.name, "r") as f:
        data = json.load(f)

    mmsis = sorted({
        str(f.get("properties", {}).get("mmsi"))
        for f in data.get("features", [])
        if f.get("properties", {}).get("mmsi") is not None
    })

    return gr.update(choices=mmsis)


# -----------------------------
# Anomaly plot
# -----------------------------
def create_anomaly_plot(df):
    df = df.sort_values("last_seen")

    fig, ax = plt.subplots(figsize=(5, 5))

    ax.plot(
        df["longitude"],
        df["latitude"],
        marker='.',
        linestyle='-',
        alpha=0.7,
        color='black'
    )

    anomal = df[df["anomaly"] == True]
    if not anomal.empty:
        ax.scatter(
            anomal["longitude"],
            anomal["latitude"],
            color="blue",
            s=60,
            label="Anomaly"
        )
        ax.legend()

    ax.set_title("Route + Anomalies")
    ax.axis('equal')
    ax.axis('off')

    return fig


# -----------------------------
# Clean route plot (UPDATED)
# -----------------------------
def create_clean_plot(df):
    df = df.sort_values("last_seen")

    fig, ax = plt.subplots(figsize=(5, 5))

    ax.plot(
        df["longitude"],
        df["latitude"],
        marker='.',
        linestyle='-',
        alpha=0.7,
        color='red'   # ✅ updated color
    )

    #ax.set_title("Route")  # ✅ removed any extra wording
    ax.axis('equal')
    ax.axis('off')

    return fig


# -----------------------------
# Main processing
# -----------------------------
def filter_vessel(geojson_file, mmsi):

    if geojson_file is None:
        raise gr.Error("Upload GeoJSON first")
    if not mmsi:
        raise gr.Error("Select MMSI")

    with open(geojson_file.name, "r") as f:
        data = json.load(f)

    rows = []

    for ftr in data.get("features", []):
        props = ftr.get("properties", {})
        if str(props.get("mmsi")) != str(mmsi):
            continue

        geom = ftr.get("geometry", {})
        coords = geom.get("coordinates")

        if not coords:
            continue

        lon, lat = coords if geom.get("type") == "Point" else (coords[0], coords[1])

        rows.append({
            "mmsi": props.get("mmsi"),
            "latitude": lat,
            "longitude": lon,
            "speed_knots": props.get("speed_knots"),
            "course_deg": props.get("course_deg"),
            "last_seen": props.get("timestamp") or props.get("last_seen"),
        })

    df = pd.DataFrame(rows)

    if df.empty:
        raise gr.Error("No data for selected MMSI")

    # -----------------------------
    # timestamp fix
    # -----------------------------
    df["last_seen"] = df["last_seen"].astype(str).str.strip()
    df["last_seen"] = parse_timestamp(df["last_seen"])
    df = df.dropna(subset=["last_seen"]).sort_values("last_seen").reset_index(drop=True)

    if df.empty:
        raise gr.Error("No valid timestamps after parsing")

    # -----------------------------
    # features
    # -----------------------------
    df["speed_knots"] = df["speed_knots"].fillna(0)
    df["speed_kmh_reported"] = df["speed_knots"] * 1.852

    df["speed_anomaly"] = False
    df["teleportation_anomaly"] = False
    df["turn_anomaly"] = False
    df["anomaly"] = False

    if len(df) > 1:

        df["prev_lat"] = df["latitude"].shift(1)
        df["prev_lon"] = df["longitude"].shift(1)
        df["prev_time"] = df["last_seen"].shift(1)
        df["prev_course"] = df["course_deg"].shift(1)

        # distance
        df["distance_km"] = df.apply(
            lambda r: haversine(r["prev_lat"], r["prev_lon"], r["latitude"], r["longitude"])
            if pd.notnull(r["prev_lat"]) else 0,
            axis=1,
        )

        # time fix
        df["time_hours"] = (df["last_seen"] - df["prev_time"]).dt.total_seconds() / 3600
        df.loc[df["time_hours"] < 0.001, "time_hours"] = np.nan

        # computed speed
        df["computed_speed_kmh"] = df["distance_km"] / df["time_hours"]
        df["computed_speed_kmh"] = df["computed_speed_kmh"].replace([np.inf, -np.inf], np.nan)
        df.loc[df["computed_speed_kmh"] > 1000, "computed_speed_kmh"] = np.nan

        # SPEED anomaly
        df["speed_anomaly"] = df["speed_kmh_reported"] > 70

        # TELEPORT anomaly
        df["teleportation_anomaly"] = df["computed_speed_kmh"] > 120

        # TURN anomaly
        df["turn_angle"] = df.apply(
            lambda r: angle_diff(r["course_deg"], r["prev_course"])
            if pd.notnull(r["prev_course"]) and pd.notnull(r["course_deg"])
            else np.nan,
            axis=1
        )

        moving = df["speed_kmh_reported"] > 5
        df["turn_anomaly"] = (df["turn_angle"] > 90) & moving

        df["anomaly"] = (
            df["speed_anomaly"] |
            df["teleportation_anomaly"] |
            df["turn_anomaly"]
        )

    # -----------------------------
    # Plotly map
    # -----------------------------
    df["label"] = df["anomaly"].map({True: "Anomaly", False: "Normal"})

    fig_map = px.scatter_mapbox(
        df,
        lat="latitude",
        lon="longitude",
        color="label",
        hover_data=[
            "speed_knots",
            "speed_kmh_reported",
            "computed_speed_kmh",
            "turn_angle",
            "last_seen"
        ],
        zoom=5,
        height=700,
    )

    fig_map.add_trace(
        dict(
            type="scattermapbox",
            mode="lines+markers",
            lat=df["latitude"],
            lon=df["longitude"],
            line=dict(width=2, color="black"),
            name="Route",
        )
    )

    fig_map.update_layout(
        mapbox_style="open-street-map",
        margin=dict(l=0, r=0, t=30, b=0),
        title=f"MMSI {mmsi} - AIS Anomaly Detection",
    )

    # -----------------------------
    # matplotlib plots
    # -----------------------------
    fig_anomaly = create_anomaly_plot(df)
    fig_clean = create_clean_plot(df)

    # -----------------------------
    # export
    # -----------------------------
    json_output = df.to_json(orient="records", indent=2, date_format="iso")

    tmp_file = tempfile.NamedTemporaryFile(suffix=".geojson", delete=False)
    with open(tmp_file.name, "w") as f:
        json.dump(data, f, indent=2)

    return (
        fig_map,
        df,
        json_output,
        fig_anomaly,
        fig_clean,
        tmp_file.name,
        "OK",
        "OK",
        "OK",
        f"Total points: {len(df)}"
    )


# -----------------------------
# UI
# -----------------------------
with gr.Blocks(title="AIS Anomaly Detection") as demo:

    gr.Markdown("# 🚢 AIS Vessel Route + Clean Plot Separation")

    with gr.Row():
        geojson_file = gr.File(label="Upload GeoJSON")
        mmsi_dropdown = gr.Dropdown(label="Select MMSI", choices=[], interactive=True)

    geojson_file.change(load_geojson, geojson_file, mmsi_dropdown)

    run_btn = gr.Button("Analyze Vessel", variant="primary")

    map_out = gr.Plot()
    table_out = gr.Dataframe()
    json_out = gr.Code(language="json")

    mpl_anomaly_out = gr.Plot(label="Route + Anomalies")
    mpl_clean_out = gr.Plot(label="Route")

    download_out = gr.File()

    speed_out = gr.Textbox(label="⚡ Speed Anomaly")
    teleport_out = gr.Textbox(label="🧭 Teleportation Anomaly")
    turn_out = gr.Textbox(label="🔄 Turn Anomaly")
    summary_out = gr.Textbox(label="📊 Summary")

    run_btn.click(
        filter_vessel,
        [geojson_file, mmsi_dropdown],
        [
            map_out,
            table_out,
            json_out,
            mpl_anomaly_out,
            mpl_clean_out,
            download_out,
            speed_out,
            teleport_out,
            turn_out,
            summary_out
        ],
    )

if __name__ == "__main__":
    ## LAunch
    demo.launch()