DSE Cal Adapt Climate Pipeline setup tutorial (w/ coiled version)

#setup check
import sys
import os

errors = []

# Check thre exist we can import key packages
try:
    import xarray
    print(f"✓ xarray {xarray.__version__}")
except ImportError:
    errors.append("xarray not installed")

try:
    import geopandas
    print(f"✓ geopandas {geopandas.__version__}")
except ImportError:
    errors.append("geopandas not installed")

try:
    import coiled
    print(f"✓ coiled installed")
except ImportError:
    errors.append("coiled not installed")

# Check library is accessible
PROJECT_ROOT = os.path.abspath(os.path.join(os.getcwd(), "..", ".."))
sys.path.insert(0, os.path.join(PROJECT_ROOT, "lib"))

try:
    from andrewAdaptLibrary import CatalogExplorer, VARIABLE_MAP
    print(f"andrewAdaptLibrary imported")
    print(f"  Variables: {list(VARIABLE_MAP.keys())}")
except ImportError as e:
    errors.append(f"andrewAdaptLibrary: {e}")

# Check mega NPS shapefile exists
nps_path = os.path.join(
    PROJECT_ROOT,
    "USA_National_Park_Service_Lands_20170930_4993375350946852027",
    "USA_Federal_Lands.shp",
)
if os.path.exists(nps_path):
    print(f"NPS shapefile found")
else:
    errors.append("NPS shapefile not found")

print()
if errors:
    print("X SETUP ISSUES:")
    for e in errors:
        print(f"   - {e}")
else:
    print(" yay! go to Part 1.")

from andrewAdaptLibrary import ParkCatalog

NPS_SHP = os.path.join(
    PROJECT_ROOT,
    "USA_National_Park_Service_Lands_20170930_4993375350946852027",
    "USA_Federal_Lands.shp",
)
catalog = ParkCatalog(NPS_SHP)

# find yosemite, then see what data is available (CA park, fast)
catalog.search("yosemite")
print()
catalog.what_is_available("Yosemite National Park")

# rocky mountain is in colorado, outside the LOCA2 grid.
# this one actually probes WRF zarr stores on S3 to confirm what exists.
catalog.search("rocky mountain")
print()
catalog.what_is_available("Rocky Mountain National Park")

# get_boundary() pulls a park's shape from the big NPS shapefile! and even returns the same kind of GeoDataFrame as load_boundary(), so
# you can pass easily throw it straight into get_climate_data().

rocky = catalog.get_boundary("Rocky Mountain National Park")
print(f"type:   {type(rocky).__name__}")
print(f"crs:    {rocky.crs}")
print(f"bounds: {rocky.total_bounds}")

# once one has thier coiled cluster running (see Part 1 below),
# fetching WRF data for rocky mountain looks like this:
#
# from andrewAdaptLibrary import get_climate_data
#
# data = get_climate_data(
#     variables=["T_Max"],
#     scenarios=["Historical Climate"],
#     boundary=rocky,
#     time_slice=(2000, 2014),
#     timescale="monthly",
#     activity="WRF",        # WRF instead of default LOCA2
#     grid="d02",            # 9km resolution (d01=45km, d02=9km)
#     backend="coiled",
#     coiled_cluster=cluster,
# )
# data["T_Max"].head()
#
# its all the same interface as LOCA2 just agotta add activity="WRF" and grid="d02"

import coiled

try:
    # This will fail if not logged in
    print("✓ Coiled is configured!")
    print(f"  View workspace: https://cloud.coiled.io")
except Exception as e:
    print(f" Not logged in: {e}")
    print("  Run 'coiled login' in your terminal first.")

!coiled env list 2>/dev/null | grep -E "(caladapt|Name)" || echo "Run 'coiled login' first"

import coiled
from dask.distributed import Client
import time
#this just starts the cluster btw... no cal_adapt stuff. Also, note that even though this takes only a few minutes, if you leave it idle (as in don't run the next cells) it'll shut down.

print("Starting test cluster...")
print("(~2 min first time, faster after)")
print()
t0 = time.perf_counter()

cluster = coiled.Cluster(
    name="tutorial-test",              # name shows up in coiled dashboard
    region="us-west1",                 # GCP region closest to S3 data
    n_workers=1,                       # just 1 worker for this test
    worker_memory="8 GiB",             # memory per worker (limited options: 2,4,8,16,32 GiB)
    idle_timeout="5 minutes",          # auto-shutdown after 5 min idle (saves cpu hours)
    package_sync=True,                 # sync local packages to workers (builds env automatically)
)

# get_client() connects  local dask to remtoe scheduler
client = cluster.get_client()
print(f"\n✓ Cluster ready in {time.perf_counter() - t0:.0f}s")
print(f"  Workers: {len(client.scheduler_info()['workers'])}")

# Run a simple task on the remote worker
import dask

@dask.delayed
def say_hello():
    import socket
    return f"Hello from {socket.gethostname()} in us-west-2!"

result = say_hello().compute()
print(result)
print("\n Coiled is working!")

# Always shut down when done- so we conserve coield credits and so i don't get charged on AWS a load  🌝
cluster.close()
print("Cluster shut down. 🌝")

import sys
import os

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

from andrewAdaptLibrary import (
    CatalogExplorer,                          # get what data exists
    load_boundary, get_lat_lon_bounds,        # work w/ shapefiles
    get_climate_data,                         # Main data fetching function
    VARIABLE_MAP, SCENARIO_MAP,               # Human-readable mappings
)

print("Library loaded!")
print(f"Variables: {list(VARIABLE_MAP.keys())}")
print(f"Scenarios: {list(SCENARIO_MAP.keys())}")

cat = CatalogExplorer()
print(cat)
print()

print("Available at 3km monthly resolution:")
print(f"  Variables: {len(cat.variables())}")
print(f"  Scenarios: {len(cat.scenarios())}")
print(f"  Total Zarr stores: {cat.catalog_size}")

# How many simulations per variable x scenario?
summary = cat.summary()

# make it prettier
summary.columns = ["Historical", "SSP 2-4.5", "SSP 3-7.0", "SSP 5-8.5"]
summary.index.name = "Variable"

# rename the variable IDs to eay to rread names
var_names = {
    "tasmax": "Max Temp",
    "tasmin": "Min Temp", 
    "pr": "Precip",
    "hursmax": "Max Humidity",
    "hursmin": "Min Humidity",
    "huss": "Specific Humidity",
    "rsds": "Solar Radiation",
    "uas": "Wind (E-W)",
    "vas": "Wind (N-S)",
    "wspeed": "Wind Speed",
}
summary.index = summary.index.map(lambda x: var_names.get(x, x))

summary.style.set_caption("Simulations available per variable × scenario")
#TODO for loca2 only, add one for WRF

# Which climate GCMs are available?
gcms = cat.gcms("tasmax", "historical")
print(f"{len(gcms)} GCMs for tasmax/historical:")
for gcm in gcms:
    print(f"  - {gcm}")

# Load Joshua Tree National Park boundary from the mega NPS shapefile
from andrewAdaptLibrary import ParkCatalog

NPS_SHP = os.path.join(
    PROJECT_ROOT,
    "USA_National_Park_Service_Lands_20170930_4993375350946852027",
    "USA_Federal_Lands.shp",
)
catalog = ParkCatalog(NPS_SHP)
boundary = catalog.get_boundary("Joshua Tree National Park")

print(f"Loaded: {len(boundary)} polygon(s)")
print(f"CRS: {boundary.crs}")

lat_bounds, lon_bounds = get_lat_lon_bounds(boundary)
print(f"Latitude:  {lat_bounds[0]:.2f}°N to {lat_bounds[1]:.2f}°N")
print(f"Longitude: {abs(lon_bounds[0]):.2f}°W to {abs(lon_bounds[1]):.2f}°W")

import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=(8, 6))
boundary.plot(ax=ax, edgecolor="black", facecolor="lightgreen", alpha=0.6)
ax.set_title("Joshua Tree National Park")
ax.set_xlabel("Longitude")
ax.set_ylabel("Latitude")
plt.show()

import coiled
import time
print("Starting cluster...")
t0 = time.perf_counter()

cluster = coiled.Cluster(
    name="tutorial-fetch",
    region="us-west1",                 # GCP region closest to S3 data in us-west-2
    n_workers=8,                       # 3 workers = fetch 3 vars in parallel
    worker_memory="16 GiB",             # each worker gets 8 GiB RAM
    spot_policy="spot_with_fallback",  # use spot instances (cheaper), fall back to on-demand if unavailable
    idle_timeout="5 minutes",          # shut down if idle (don't waste credits)
    package_sync=True,                 # sync local packages to workers (builds env automatically)
)

# this connects local python session to the remote scheduler
# all dask operations now run on the cluster, not laptop!
client = cluster.get_client()
print(f"\n✓ Ready in {time.perf_counter() - t0:.0f}s")
print(f"  Dashboard: {client.dashboard_link}")

# Fetch climate data using my library wrapper around calAdapt
# this gives work to the remote workers, they do the heavy lifting
print("Fetching data...")
t0 = time.perf_counter()

data = get_climate_data(
    variables=["T_Max", "T_Min", "Precip"],  # human-readable names (mapped to tasmax, tasmin, pr)
    scenarios=["Historical Climate"],         # mapped to "historical" experiment_id, all four would be: scenarios=["Historical Climate", "SSP 2-4.5", "SSP 3-7.0", "SSP 5-8.5"]
    boundary=boundary,                        # GeoDataFrame gets serialized to WKT automatically
    time_slice=(2010, 2014),                  # 5 years for demo-- keep it small to save credits
    timescale="monthly",                      # "monthly", "daily", or "yearly" (monthly is default)
    backend="coiled",                         # use Coiled cluster (vs "direct_s3" for local)
    coiled_cluster=cluster,                   # pass the cluster we just created
)

# data is a dict: {variable: DataFrame} where each DataFrame has all scenarios
print(f"\n✓ Done in {time.perf_counter() - t0:.1f}s")
print(f"  Got {len(data)} variables")

# What did we get back? Let's peek at the structure

print("=== What `data` looks like ===")
print(f"Type: {type(data)}")
print(f"Keys: {list(data.keys())}")
print()

# Each key is a variable name, each value is a DataFrame with ALL scenarios
for var, df in data.items():
    print(f'data["{var}"] → DataFrame')
    print(f"    Shape: {df.shape[0]} rows × {df.shape[1]} cols")
    print(f"    Columns: {list(df.columns)}")
    print(f"    Scenarios in df: {df['scenario'].unique().tolist()}")
    print()

# Let's look inside one of those DataFrames
sample_df = data["T_Max"]

print("=== Inside the T_Max DataFrame ===")
print(f"Rows: {len(sample_df):,} (= 70 simulations × 60 months)")
print(f"Columns: {list(sample_df.columns)}")
print()
print("First 10 rows:")
sample_df.head(10)

# Always shut down the cluster!!!!
cluster.close()
print("Cluster shut down.")

# Using the data we fetched in Part 2, let's look at T_Max
# (If you skipped Part 2, you'll need to run those cells first)

# The dict key is just the variable name now!
tmax_df = data["T_Max"]

print("=== DataFrame from get_climate_data() ===")
print(f"Type: {type(tmax_df)}")
print(f"Shape: {tmax_df.shape[0]} rows × {tmax_df.shape[1]} columns")
print(f"Columns: {list(tmax_df.columns)}")
print()
print("First 10 rows:")
tmax_df.head(10)

# Each row = one month × one simulation
# print to see what's in each column

print("=== Column breakdown ===")
print(f"time: {tmax_df['time'].dtype}")
print(f"  Range: {tmax_df['time'].min()} to {tmax_df['time'].max()}")
print()
print(f"simulation: {tmax_df['simulation'].dtype}")
print(f"  Unique values: {tmax_df['simulation'].nunique()}")
print(f"  Examples: {tmax_df['simulation'].unique()[:3].tolist()}")
print()
print(f"scenario: {tmax_df['scenario'].dtype}")
print(f"  Unique values: {tmax_df['scenario'].unique().tolist()}")
print()
print(f"T_Max: {tmax_df['T_Max'].dtype}")
print(f"  Range: {tmax_df['T_Max'].min():.1f}°C to {tmax_df['T_Max'].max():.1f}°C")
print(f"  Mean: {tmax_df['T_Max'].mean():.1f}°C")

# s3_paths returns the actual S3 bucket locations
paths = cat.s3_paths("tasmax", "historical")
print(f"{len(paths)} Zarr stores for tasmax/historical")
print()
print("First 3:")
for p in paths[:3]:
    print(f"  {p['source_id']:15s} → {p['path']}")

from andrewAdaptLibrary import boundary_to_wkt, boundary_from_wkt

# Convert to strings (this is what we send to remote workers)
wkt_list, crs_str = boundary_to_wkt(boundary)
print(f"CRS: {crs_str}")
print(f"WKT (first 150 chars): {wkt_list[0][:150]}...")

# Reconstruct on the other side (this happens on the remote worker)
reconstructed = boundary_from_wkt(wkt_list, crs_str)

# verify it worked - compare geometry bounds instead of area (area calc warns about geographic CRS)
orig_bounds = boundary.total_bounds
recon_bounds = reconstructed.total_bounds
print(f"\nRound-trip OK: {all(abs(o - r) < 1e-6 for o, r in zip(orig_bounds, recon_bounds))}")

from andrewAdaptLibrary import fetch_direct_s3

# Open Zarr stores (lazy — no download yet)
lazy = fetch_direct_s3(
    var_key="T_Max",
    experiment="historical",
    time_slice=(2010, 2010),  # Just 1 year
    lat_bounds=lat_bounds,
    lon_bounds=lon_bounds,
    catalog=cat,
)

print(f"Shape: {lazy.shape}")
print(f"Dims: {lazy.dims}")
print(f"Size if loaded: {lazy.nbytes / 1e6:.1f} MB")
print()
print("No data downloaded yet! Call .load() to fetch.")

# Example: looking up a park and running the full pipeline

def analyze_park(park_name, variables, scenarios, years, timescale="monthly", n_workers=3):
    """
    Template function for analyzing any park/region.
    
    Args:
        park_name: exact NPS unit name (use catalog.search() to find it)
        variables: list of ["T_Max", "T_Min", "Precip"]
        scenarios: list of ["Historical Climate", "SSP 2-4.5", "SSP 3-7.0", "SSP 5-8.5"]
        years: tuple like (2010, 2020)
        timescale: "monthly", "daily", or "yearly" (default "monthly")
        n_workers: how many parallel workers (default 3)
    """
    import coiled
    
    # 1. Load the boundary from the mega NPS shapefile
    boundary = catalog.get_boundary(park_name)
    print(f"Loaded boundary for {park_name}: {len(boundary)} polygon(s)")
    
    # check it's in CA roughly
    lat_bounds, lon_bounds = get_lat_lon_bounds(boundary)
    if lat_bounds[0] < 32 or lat_bounds[1] > 42:
        print("Warning: latitude outside CA range (32-42°N)")
    if lon_bounds[0] < -125 or lon_bounds[1] > -114:
        print("Warning: longitude outside CA range (125-114°W)")
    
    # 2. Start cluster
    cluster = coiled.Cluster(
        name="custom-park-analysis",
        region="us-west1",
        n_workers=n_workers,
        worker_memory="8 GiB",
        spot_policy="spot_with_fallback",
        idle_timeout="10 minutes",
        package_sync=True,             # sync local packages to workers
    )
    client = cluster.get_client()
    
    try:
        # 3. Fetch data
        data = get_climate_data(
            variables=variables,
            scenarios=scenarios,
            boundary=boundary,
            time_slice=years,
            timescale=timescale,
            backend="coiled",
            coiled_cluster=cluster,
        )
        return data
    finally:
        # 4. Always shut down
        cluster.close()
        print("Cluster shut down.")

#  (don't run this— just an example):
# data = analyze_park(
#     park_name="Death Valley National Park",
#     variables=["T_Max", "T_Min"],
#     scenarios=["Historical Climate", "SSP 5-8.5"],
#     years=(2015, 2025),
#     timescale="monthly",  # or "daily" for extreme event analysis
# )

import coiled

# Fill in the blanks
cluster = coiled.Cluster(
    name="exercise-cluster",
    region="____",             # hint: GCP region closest to the S3 data
    n_workers=2,
    worker_memory="8 GiB",     # this is filled in since coiled gets mad if you choose non-standard sizes
    idle_timeout="____ minutes",  # hint: don't waste credits
    package_sync=True,         # this syncs your local packages to workers
)
client = cluster.get_client()
print(f"Dashboard: {client.dashboard_link}")

# Check: if the cluster started, you're good!
print(f"Workers: {len(client.scheduler_info()['workers'])}")
print("Exercise 1 done!")

# SOLUTION
import coiled

cluster = coiled.Cluster(
    name="exercise-cluster",
    region="us-west1",            # GCP region closest to S3 data
    n_workers=2,
    worker_memory="8 GiB",
    idle_timeout="5 minutes",     # any reasonable timeout works
    package_sync=True,            # sync local packages to workers
)
client = cluster.get_client()
print(f"Dashboard: {client.dashboard_link}")

# Load boundary from the mega NPS shapefile
jt_boundary = catalog.get_boundary("Joshua Tree National Park")

# Fill in the blanks
data = get_climate_data(
    variables=["____"],              # hint: max temperature
    scenarios=["____"],              # hint: past observations
    boundary=____,                   # hint: variable we just loaded
    time_slice=(2015, ____),         # hint: 5 years ending in 2019
    backend="____",                  # hint: use the cloud
    coiled_cluster=____,             # hint: the cluster we started above
)

# Check your answer
def check_ex2(data):
    assert "T_Max" in data, f"Expected 'T_Max' key, got {list(data.keys())}"
    df = data["T_Max"]
    assert len(df) > 0, "DataFrame is empty"
    years = df['time'].dt.year.unique()
    assert 2015 in years and 2019 in years, f"Wrong years: {years}"
    print(f"Exercise 2 done. {len(df)} rows fetched.")

check_ex2(data)

# SOLUTION
jt_boundary = catalog.get_boundary("Joshua Tree National Park")

data = get_climate_data(
    variables=["T_Max"],
    scenarios=["Historical Climate"],
    boundary=jt_boundary,
    time_slice=(2015, 2019),
    timescale="monthly",
    backend="coiled",
    coiled_cluster=cluster,
)

print(f"Got {len(data)} result(s)")

# grab one of the dataframes we fetched earlier
# data is a dict like: {"T_Max": DataFrame, "T_Min": DataFrame, ...}
sample_df = data["T_Max"]

print("=== DataFrame Structure ===")
print(f"Shape: {sample_df.shape[0]} rows × {sample_df.shape[1]} columns")
print(f"\nColumns: {list(sample_df.columns)}")
print(f"\nData types:")
print(sample_df.dtypes)
print(f"\n=== First 5 rows ===")
sample_df.head()

# quick stats
print("=== Quick Stats ===")
print(f"Time range: {sample_df['time'].min()} to {sample_df['time'].max()}")
print(f"Unique simulations: {sample_df['simulation'].nunique()}")
print(f"\nT_Max summary (°C):")
print(sample_df['T_Max'].describe())

# Option 1: Save each variable separately
import os
output_dir = os.path.join(PROJECT_ROOT, "data", "processed")
os.makedirs(output_dir, exist_ok=True)

for var, df in data.items():
    # clean filename: "T_Max.csv"
    filename = f"{var}.csv"
    filepath = os.path.join(output_dir, filename)
    df.to_csv(filepath, index=False)
    print(f"Saved: {filename} ({len(df)} rows)")

print(f"\nFiles saved to: {output_dir}")

# Option 2: Combine all variables into one tidy dataframe
import pandas as pd

# Each DataFrame already has 'scenario' column, just add 'variable' and rename value col
tidy_dfs = []
for var, df in data.items():
    df_copy = df.copy()
    df_copy['variable'] = var
    df_copy = df_copy.rename(columns={var: 'value'})  # generic 'value' column
    tidy_dfs.append(df_copy)

tidy_df = pd.concat(tidy_dfs, ignore_index=True)
print(f"Combined shape: {tidy_df.shape}")
print(f"Columns: {list(tidy_df.columns)}")
tidy_df.head()

# save the tidy version
tidy_path = os.path.join(output_dir, "climate_data_tidy.csv")
tidy_df.to_csv(tidy_path, index=False)
print(f"Saved: {tidy_path}")

# === Minimal ex: Fetch 5 years of climate data for any park ===
import os, coiled
from andrewAdaptLibrary import ParkCatalog, get_climate_data

PROJECT_ROOT = os.path.abspath(os.path.join(os.getcwd(), "..", ".."))
NPS_SHP = os.path.join(
    PROJECT_ROOT,
    "USA_National_Park_Service_Lands_20170930_4993375350946852027",
    "USA_Federal_Lands.shp",
)

# 1. Look up your park in the mega NPS shapefile → returns GeoDataFrame
catalog = ParkCatalog(NPS_SHP)
boundary = catalog.get_boundary("Joshua Tree National Park")

# 2. Spin up cluster → returns Cluster object (workers now running in GCP us-west1)
cluster = coiled.Cluster(
    name="quick-fetch",
    region="us-west1",
    n_workers=3,
    worker_memory="8 GiB",
    spot_policy="spot_with_fallback",
    idle_timeout="10 minutes",
    package_sync=True,
)

# 3. Fetch data → returns {variable_name: DataFrame}
data = get_climate_data(
    variables=["T_Max"],
    scenarios=["Historical Climate"],
    boundary=boundary,
    time_slice=(2010, 2014),
    timescale="monthly",
    backend="coiled",
    coiled_cluster=cluster,
)

# 4. Shut down cluster
cluster.close()

# 5. Use the data
df = data["T_Max"]
print(df.head())