from pathlib import Path
from urllib.parse import urlencode

import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from datetime import datetime, timezone

from autark import *


def usgs_csv_url(days: int, min_magnitude: float) -> str:
    params = {
        "format": "csv",
        "orderby": "time-asc",
        "starttime": f"now-{days}days",
        "endtime": "now",
        "minmagnitude": str(min_magnitude),
        "limit": "20000",
    }
    return "https://earthquake.usgs.gov/fdsnws/event/1/query?" + urlencode(params)


def parse_usgs_time(value):
    if isinstance(value, datetime):
        return value
    if isinstance(value, str):
        return datetime.fromisoformat(value.replace("Z", "+00:00"))
    if isinstance(value, (int, float)):
        v = float(value)
        av = abs(v)
        # Handle epoch values in ns/us/ms/s.
        if av > 1e17:
            v = v / 1_000_000_000
        elif av > 1e14:
            v = v / 1_000_000
        elif av > 1e11:
            v = v / 1_000
        return datetime.fromtimestamp(v, tz=timezone.utc)
    raise TypeError(f"unsupported time value: {value!r}")


def main() -> None:
    url = usgs_csv_url(days=30, min_magnitude=1.0)
    frame = OnceFrame(CsvReader(url))
    print(frame.describe())

    p = P()

    p.columns(["time", "mag"]).slice(0, 400).alias(
        "mag_timeseries", frame.schema_of_columns(["time", "mag"])
    )
    p.columns("mag").reduce("mean").alias("mean_mag", frame.schema_of_columns(["mag"]))
    p.columns("depth").reduce("mean").alias(
        "mean_depth", frame.schema_of_columns(["depth"])
    )
    p.columns("mag").reduce("count").alias("row_count", frame.schema_of_columns(["mag"]))
    agg = p.columns("mag").group_reduce(
        p.columns("nst"), kind=["count", "mean", "max"]
    )
    agg = agg.order_by(agg.columns(0), ascending=False)
    agg.alias(
        "mag_by_nst",
        frame.schema_of_columns(["nst"])
        .column_f64("mag_count")
        .column_f64("mag_mean")
        .column_f64("mag_max"),
    )

    left = p.columns(None).slice(50, 550)
    right = p.columns(None).slice(0, 500)

    left.join(
        right,
        left_on=left.columns("time"),
        right_on=right.columns("time"),
        kind="asof_backward",
    ).alias("asof_backward")

    left.join(
        right,
        left_on=left.columns("time"),
        right_on=right.columns("time"),
        kind="asof_forward",
    ).alias("asof_forward")

    realized = frame.realize(p)
    print("keys:", realized.keys())
    print(realized)

    series = realized["mag_timeseries"]
    t = [parse_usgs_time(x) for x in series["time"]]
    m = series["mag"]

    fig, axes = plt.subplots(2, 1, figsize=(11, 7))

    ax = axes[0]
    ax.plot(t, m, color="#0ea5e9", linewidth=1.5)
    ax.scatter(t, m, s=10, color="#0ea5e9", alpha=0.35)
    top_idx = sorted(range(len(m)), key=lambda i: m[i], reverse=True)[:8]
    ax.scatter([t[i] for i in top_idx], [m[i] for i in top_idx], s=36, color="#f97316")
    ax.xaxis.set_major_formatter(mdates.DateFormatter("%b %d"))
    ax.set_title("Earth Pulse: Magnitude Through Time")
    ax.set_xlabel("UTC Date")
    ax.set_ylabel("Magnitude")
    ax.grid(alpha=0.25)
    ax.tick_params(axis="x", rotation=20)

    mean_mag = float(realized["mean_mag"]["mag"][0])
    mean_depth = float(realized["mean_depth"]["depth"][0])
    ax2 = axes[1]
    bars = ax2.bar(
        ["Mean Magnitude", "Mean Depth (km)"],
        [mean_mag, mean_depth],
        color=["#0ea5e9", "#fb923c"],
    )
    ax2.grid(axis="y", alpha=0.25)
    ax2.set_title("Aggregate Metrics")
    for bar in bars:
        value = bar.get_height()
        ax2.text(
            bar.get_x() + bar.get_width() / 2,
            value,
            f"{value:.2f}",
            ha="center",
            va="bottom",
        )

    fig.suptitle("USGS Earthquake Activity Snapshot")
    fig.tight_layout(rect=[0, 0, 1, 0.97])
    out_path = Path(__file__).with_name("usgs_a_plot.png")
    fig.savefig(out_path, dpi=140)
    print(f"saved plot: {out_path}")


if __name__ == "__main__":
    main()
