Full Climate Data Extraction: Joshua Tree & Mojave

import sys
import os
import time
import pandas as pd
import coiled

PROJECT_ROOT = os.path.abspath(os.path.join(os.getcwd(), "..", ".."))
sys.path.insert(0, os.path.join(PROJECT_ROOT, "lib"))

from andrewAdaptLibrary import (
    ParkCatalog,
    get_lat_lon_bounds,
    get_climate_data,
    VARIABLE_MAP,
    SCENARIO_MAP,
)

# Output directory
OUTPUT_DIR = os.path.join(PROJECT_ROOT, "data", "csv", "full_extraction")
os.makedirs(OUTPUT_DIR, exist_ok=True)

# Load the mega NPS shapefile for park lookups
NPS_SHP = os.path.join(
    PROJECT_ROOT,
    "USA_National_Park_Service_Lands_20170930_4993375350946852027",
    "USA_Federal_Lands.shp",
)
catalog = ParkCatalog(NPS_SHP)

print(f"Variables: {list(VARIABLE_MAP.keys())}")
print(f"Scenarios: {list(SCENARIO_MAP.keys())}")
print(f"Output dir: {OUTPUT_DIR}")

parks = ["Joshua Tree National Park", "Mojave National Preserve"]

boundaries = {}
for park_name in parks:
    boundaries[park_name] = catalog.get_boundary(park_name)
    lat, lon = get_lat_lon_bounds(boundaries[park_name])
    print(f"{park_name}: {lat[0]:.2f}-{lat[1]:.2f}°N, {abs(lon[0]):.2f}-{abs(lon[1]):.2f}°W")

print("Starting cluster on GCP...")
t_cluster_start = time.perf_counter()

cluster = coiled.Cluster(
    name="full-extraction",
    region="us-west1",
    n_workers=32,
    worker_vm_types=["n1-highmem-4"],  # 4 vCPUs, 26 GiB RAM per worker
    spot_policy="spot_with_fallback",
    idle_timeout="70 minutes",         # longer timeout so cluster survives between cells
    package_sync=True,
)

client = cluster.get_client()
t_cluster_ready = time.perf_counter()

print(f"\n✓ Cluster ready in {t_cluster_ready - t_cluster_start:.0f}s")
print(f"  Workers: {len(client.scheduler_info()['workers'])}")
print(f"  Dashboard: {client.dashboard_link}")

variables = ["T_Max", "T_Min"]
historical_scenarios = ["Historical Climate"]
future_scenarios = ["SSP 2-4.5", "SSP 3-7.0", "SSP 5-8.5"]

all_timings = {}
t_total_start = time.perf_counter()

for park_name, boundary in boundaries.items():
    # Clean name for filenames: "Joshua Tree National Park" -> "JoshuaTree"
    short_name = park_name.replace("National Park", "").replace("National Preserve", "").strip().replace(" ", "")

    print(f"\n{'='*60}")
    print(f"  {park_name}")
    print(f"{'='*60}")
    
    park_dfs = {var: [] for var in variables}
    
    # --- Historical: 1950-2014 ---
    print(f"\n  gettin histproal (1950-2014)...")
    t0 = time.perf_counter()
    
    hist_data = get_climate_data(
        variables=variables,
        scenarios=historical_scenarios,
        boundary=boundary,
        time_slice=(1950, 2014),
        timescale="monthly",
        backend="coiled",
        coiled_cluster=cluster,
    )
    
    t_hist = time.perf_counter() - t0
    all_timings[f"{short_name}_historical"] = t_hist
    
    for var in variables:
        df = hist_data[var]
        print(f"    {var}: {len(df):,} rows")
        park_dfs[var].append(df)
    print(f"  Historical done in {t_hist:.1f}s")
    
    # --- Future SSPs: 2015-2100 ---
    print(f"\n  getting all SSP scens (2015-2100)...")
    t0 = time.perf_counter()
    
    future_data = get_climate_data(
        variables=variables,
        scenarios=future_scenarios,
        boundary=boundary,
        time_slice=(2015, 2100),
        timescale="monthly",
        backend="coiled",
        coiled_cluster=cluster,
    )
    
    t_future = time.perf_counter() - t0
    all_timings[f"{short_name}_future"] = t_future
    
    for var in variables:
        df = future_data[var]
        print(f"    {var}: {len(df):,} rows")
        park_dfs[var].append(df)
    print(f"  Future scenarios done in {t_future:.1f}s")
    
    # --- Join & save ---
    print(f"\n  Savng CSVs...")
    for var in variables:
        combined = pd.concat(park_dfs[var], ignore_index=True)
        combined["park"] = short_name
        filename = f"{short_name}_{var}.csv"
        filepath = os.path.join(OUTPUT_DIR, filename)
        combined.to_csv(filepath, index=False)
        print(f"    Saved: {filename} ({len(combined):,} rows)")

t_total = time.perf_counter() - t_total_start
print(f"\n{'='*60}")
print(f"  ALL DATA FETCHED")
print(f"{'='*60}")
print(f"Total fetch time: {t_total:.1f}s ({t_total/60:.1f} min)")

cluster.close()
print("Cluster shut down.")

print("=== Timing Summary ===")
print(f"Cluster spin-up: {t_cluster_ready - t_cluster_start:.0f}s")
print()
for label, t in all_timings.items():
    print(f"{label}: {t:.1f}s ({t/60:.1f} min)")
print()
print(f"Total data fetch: {t_total:.1f}s ({t_total/60:.1f} min)")
print(f"Total including cluster: {t_cluster_ready - t_cluster_start + t_total:.0f}s ({(t_cluster_ready - t_cluster_start + t_total)/60:.1f} min)")

print("=== Output Files ===")
for f in sorted(os.listdir(OUTPUT_DIR)):
    if f.endswith(".csv"):
        path = os.path.join(OUTPUT_DIR, f)
        size_mb = os.path.getsize(path) / 1e6
        df = pd.read_csv(path, nrows=5)
        row_count = sum(1 for _ in open(path)) - 1  # minus header
        print(f"\n{f}")
        print(f"  Size: {size_mb:.1f} MB")
        print(f"  Rows: {row_count:,}")
        print(f"  Columns: {list(df.columns)}")
        print(f"  Scenarios: {pd.read_csv(path, usecols=['scenario'])['scenario'].unique().tolist()}")

# Quick peek at one file
sample = pd.read_csv(os.path.join(OUTPUT_DIR, "JoshuaTree_T_Max.csv"))
print(f"Shape: {sample.shape}")
print(f"Years: {pd.to_datetime(sample['time']).dt.year.min()} to {pd.to_datetime(sample['time']).dt.year.max()}")
print(f"Scenarios: {sample['scenario'].unique().tolist()}")
print(f"Simulations: {sample['simulation'].nunique()}")
print()
sample.head()

for park_name in boundaries:
    short_name = park_name.replace("National Park", "").replace("National Preserve", "").strip().replace(" ", "")
    tmax = pd.read_csv(os.path.join(OUTPUT_DIR, f"{short_name}_T_Max.csv"))
    tmin = pd.read_csv(os.path.join(OUTPUT_DIR, f"{short_name}_T_Min.csv"))
    
    # Merge on shared columns to align rows
    merge_cols = ["simulation", "time", "scenario", "timescale", "park"]
    merged = tmax.merge(tmin[merge_cols + ["T_Min"]], on=merge_cols)
    merged["T_Avg"] = (merged["T_Max"] + merged["T_Min"]) / 2
    
    # Save
    out_path = os.path.join(OUTPUT_DIR, f"{short_name}_T_Avg.csv")
    merged.to_csv(out_path, index=False)
    print(f"{short_name}_T_Avg.csv: {len(merged):,} rows, T_Avg range: {merged['T_Avg'].min():.1f}°C to {merged['T_Avg'].max():.1f}°C")