import os, json, time, threading, io
import numpy as np
import torch
import torch.nn as nn
import requests
from PIL import Image
from matplotlib.colors import LinearSegmentedColormap
from huggingface_hub import hf_hub_download, HfApi, create_repo
from pysteps.motion.lucaskanade import dense_lucaskanade
from pysteps.extrapolation.semilagrangian import extrapolate
from pyproj import Proj
from datetime import datetime, timedelta
import gradio as gr
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse

HF_TOKEN=os.environ.get("HF_TOKEN","")
MODEL_REPO="NovatasticRoScript/himawari-nowcast"
DATASET_REPO="NovatasticRoScript/himawari-live-cache"
SEQ=9; PRED=18; LK=4; W=H=640; IN_CH=6
DISPLAY_BBOX=(103.,-3.,139.,35.)
PAR_POLY=[(115.,5.),(115.,15.),(120.,21.),(120.,25.),(135.,25.),(135.,5.)]
SLIDER="https://slider.cira.colostate.edu"
ZOOM=3
FRAMES_DIR="/tmp/frames"
UPDATE_INTERVAL_SEC=600

os.makedirs(FRAMES_DIR, exist_ok=True)

# =============================================================================
# Himawari full-disk geostationary projection
# =============================================================================
_HIMA_PROJ = Proj(proj='geos', h=35785863.0, lon_0=140.7, a=6378137.0, b=6356752.3, sweep='x')
_FULL_DISK_EXTENT = 5500000.035308

def _lonlat_to_frac(lon, lat):
    x, y = _HIMA_PROJ(lon, lat)
    fx = (x + _FULL_DISK_EXTENT) / (2*_FULL_DISK_EXTENT)
    fy = (_FULL_DISK_EXTENT - y) / (2*_FULL_DISK_EXTENT)
    return fx, fy

def _bbox_frac_range():
    lon_min, lat_min, lon_max, lat_max = DISPLAY_BBOX
    corners = [(lon_min,lat_min),(lon_min,lat_max),(lon_max,lat_min),(lon_max,lat_max)]
    fxs, fys = [], []
    for lon, lat in corners:
        fx, fy = _lonlat_to_frac(lon, lat)
        fxs.append(fx); fys.append(fy)
    return min(fxs), max(fxs), min(fys), max(fys)

def _crop_resize_to_region(full_disk_arr):
    h, w = full_disk_arr.shape
    fx0, fx1, fy0, fy1 = _bbox_frac_range()
    x0, x1 = int(fx0*w), int(fx1*w)
    y0, y1 = int(fy0*h), int(fy1*h)
    x0,x1 = max(0,x0), min(w,x1); y0,y1 = max(0,y0), min(h,y1)
    crop = full_disk_arr[y0:y1, x0:x1] if (x1>x0 and y1>y0) else full_disk_arr
    img = Image.fromarray((np.clip(crop,0,1)*255).astype(np.uint8)).resize((W,H), Image.Resampling.BILINEAR)
    return np.array(img, dtype=np.float32)/255.

def _bbox_tile_range(zoom):
    n = 2**zoom
    fx0, fx1, fy0, fy1 = _bbox_frac_range()
    col0,col1 = max(0,int(fx0*n)), min(n-1,int(fx1*n))
    row0,row1 = max(0,int(fy0*n)), min(n-1,int(fy1*n))
    return row0,row1,col0,col1,n

# =============================================================================
# Colormap + model
# =============================================================================
ir_colors=[(0.00,(0.05,0.05,0.05)),(0.30,(0.15,0.20,0.35)),(0.50,(0.00,0.65,0.90)),
           (0.65,(0.00,0.75,0.00)),(0.80,(1.00,0.85,0.00)),(0.92,(0.90,0.10,0.00)),
           (1.00,(1.00,1.00,1.00))]
cmap=LinearSegmentedColormap.from_list("ir1",ir_colors,N=256)

class RCN(nn.Module):
    def __init__(self):
        super().__init__()
        self.enc=nn.Sequential(nn.Conv2d(IN_CH,32,3,padding=1),nn.BatchNorm2d(32),nn.LeakyReLU(0.2,inplace=True),
                               nn.Conv2d(32,64,3,padding=2,dilation=2),nn.BatchNorm2d(64),nn.LeakyReLU(0.2,inplace=True))
        self.ref=nn.Sequential(nn.Conv2d(64,64,3,padding=1),nn.BatchNorm2d(64),nn.ReLU(inplace=True),
                               nn.Conv2d(64,32,3,padding=1),nn.BatchNorm2d(32),nn.ReLU(inplace=True))
        self.out=nn.Sequential(nn.Conv2d(32,1,1),nn.Sigmoid())
    def forward(self,b13,b08,b03,prior,trend,coast):
        x=torch.cat([b13,b08,b03,prior,trend,coast],dim=1)
        r=self.out(self.ref(self.enc(x)))
        return torch.clamp(prior+(r-0.5)*0.2,0.,1.)

model=RCN()
try:
    p=hf_hub_download(repo_id=MODEL_REPO,filename="model.pt",token=HF_TOKEN or None)
    model.load_state_dict(torch.load(p,map_location="cpu"))
    print("Model loaded.")
except Exception as e:
    print(f"Model load failed: {e}")
model.eval()

zeros_t=torch.zeros(1,1,H,W); coast_t=torch.zeros(1,1,H,W)

# =============================================================================
# SLIDER fetch
# =============================================================================
def slider_ts(product="band_13"):
    url=f"{SLIDER}/data/json/himawari/full_disk/{product}/latest_times.json"
    r=requests.get(url,timeout=10)
    r.raise_for_status()
    d=r.json()
    if isinstance(d,list):
        return str(d[0])
    for k in ("timestamps_int","timestamps","times"):
        if k in d:
            return str(d[k][0])
    raise RuntimeError(f"Unknown SLIDER JSON shape: {list(d.keys()) if isinstance(d,dict) else type(d)}")

def fetch_slider_raw(ts, product="band_13"):
    date_path=f"{ts[:4]}/{ts[4:6]}/{ts[6:8]}"
    url=f"{SLIDER}/data/imagery/{date_path}/himawari---full_disk/{product}/{ts}/00/000_000.png"
    r=requests.get(url,timeout=10)
    r.raise_for_status()
    img=Image.open(io.BytesIO(r.content)).convert("L")
    return np.array(img,dtype=np.float32)/255.

def fetch_slider_region(ts, zoom=ZOOM, product="band_13"):
    row0,row1,col0,col1,n = _bbox_tile_range(zoom)
    date_path=f"{ts[:4]}/{ts[4:6]}/{ts[6:8]}"
    tiles={}
    tile_h=tile_w=None
    for row in range(row0,row1+1):
        for col in range(col0,col1+1):
            url=f"{SLIDER}/data/imagery/{date_path}/himawari---full_disk/{product}/{ts}/{zoom:02d}/{row:03d}_{col:03d}.png"
            try:
                r=requests.get(url,timeout=10)
                r.raise_for_status()
                img=Image.open(io.BytesIO(r.content)).convert("L")
                arr=np.array(img,dtype=np.float32)/255.
                tiles[(row,col)]=arr
                if tile_h is None: tile_h,tile_w=arr.shape
            except Exception as e:
                print(f"TILE z{zoom} {row:03d}_{col:03d} FAILED: {e}")
    if not tiles:
        raise RuntimeError(f"No tiles fetched at zoom {zoom}")

    n_rows=row1-row0+1; n_cols=col1-col0+1
    mosaic=np.zeros((tile_h*n_rows,tile_w*n_cols),dtype=np.float32)
    for (row,col),arr in tiles.items():
        r_off=(row-row0)*tile_h; c_off=(col-col0)*tile_w
        mosaic[r_off:r_off+tile_h,c_off:c_off+tile_w]=arr

    fx0,fx1,fy0,fy1 = _bbox_frac_range()
    mosaic_fx0,mosaic_fx1 = col0/n,(col1+1)/n
    mosaic_fy0,mosaic_fy1 = row0/n,(row1+1)/n
    mh,mw = mosaic.shape
    px0=int((fx0-mosaic_fx0)/(mosaic_fx1-mosaic_fx0)*mw)
    px1=int((fx1-mosaic_fx0)/(mosaic_fx1-mosaic_fx0)*mw)
    py0=int((fy0-mosaic_fy0)/(mosaic_fy1-mosaic_fy0)*mh)
    py1=int((fy1-mosaic_fy0)/(mosaic_fy1-mosaic_fy0)*mh)
    px0,px1=max(0,px0),min(mw,px1); py0,py1=max(0,py0),min(mh,py1)
    crop = mosaic[py0:py1,px0:px1] if (px1>px0 and py1>py0) else mosaic
    img = Image.fromarray((np.clip(crop,0,1)*255).astype(np.uint8)).resize((W,H), Image.Resampling.BILINEAR)
    return np.array(img, dtype=np.float32)/255.

def fetch_slider(ts, product="band_13"):
    try:
        return fetch_slider_region(ts, zoom=ZOOM, product=product)
    except Exception as e:
        print(f"Zoom {ZOOM} tile mosaic failed ({e}) -- falling back to full-disk zoom00")
        raw = fetch_slider_raw(ts, product)
        return _crop_resize_to_region(raw)

def build_live_seq():
    ts=slider_ts()
    base=datetime.strptime(ts,"%Y%m%d%H%M%S")
    frames=[]
    for i in range(SEQ-1,-1,-1):
        stamp=(base-timedelta(minutes=10*i)).strftime("%Y%m%d%H%M%S")
        try:
            frames.append(fetch_slider(stamp))
        except Exception as e:
            print(f"FETCH_SLIDER FAILED for {stamp}: {e}")
            frames.append(frames[-1] if frames else np.zeros((H,W),dtype=np.float32))
    while len(frames)<SEQ:
        frames.insert(0,frames[0])
    return np.array(frames[-SEQ:]),base

def save_frame_png(data, path):
    rgb = (cmap(np.clip(data,0,1))[...,:3]*255).astype(np.uint8)
    Image.fromarray(rgb).save(path)

# =============================================================================
# Background update loop
# =============================================================================
def update_loop():
    while True:
        try:
            history, base = build_live_seq()
            v = dense_lucaskanade(history[-LK:], verbose=False)
            hor = np.array(extrapolate(history[-1], v, timesteps=PRED, outval="min"))

            curr = torch.tensor(history[-1], dtype=torch.float32).unsqueeze(0).unsqueeze(0)
            last = torch.tensor(history[-2], dtype=torch.float32).unsqueeze(0).unsqueeze(0)
            preds=[]
            with torch.no_grad():
                for i in range(PRED):
                    prior = torch.tensor(hor[i], dtype=torch.float32).unsqueeze(0).unsqueeze(0)
                    ref = model(curr, zeros_t, zeros_t, prior, curr-last, coast_t)
                    preds.append(ref.squeeze().numpy()); last=curr; curr=ref

            for i in range(SEQ):
                save_frame_png(history[i], os.path.join(FRAMES_DIR, f"obs_{i}.png"))
            for i in range(PRED):
                save_frame_png(preds[i], os.path.join(FRAMES_DIR, f"fcst_{i}.png"))

            manifest = {
                "ready": True,
                "base_time": base.strftime("%Y-%m-%dT%H:%M:%SZ"),
                "obs_count": SEQ, "fcst_count": PRED,
                "obs_step_min": 10, "fcst_step_min": 10,
                "version": int(time.time()),
            }
            with open(os.path.join(FRAMES_DIR,"manifest.json"),"w") as f:
                json.dump(manifest, f)
            print(f"✅ Updated frames @ {base} UTC")

            if HF_TOKEN:
                try:
                    api = HfApi(token=HF_TOKEN)
                    create_repo(DATASET_REPO, repo_type="dataset", token=HF_TOKEN, exist_ok=True)
                    fname = f"live_{base:%Y%m%d_%H%M}.npy"
                    npy_path = f"/tmp/{fname}"
                    np.save(npy_path, history[-1])
                    api.upload_file(path_or_fileobj=npy_path, path_in_repo=f"frames/{fname}",
                                     repo_id=DATASET_REPO, repo_type="dataset", token=HF_TOKEN)
                except Exception as e:
                    print(f"Dataset push failed: {e}")

        except Exception as e:
            print(f"⚠️ Update loop error: {e}")

        time.sleep(UPDATE_INTERVAL_SEC)

if not os.path.exists(os.path.join(FRAMES_DIR,"manifest.json")):
    with open(os.path.join(FRAMES_DIR,"manifest.json"),"w") as f:
        json.dump({"ready": False, "version": 0}, f)

threading.Thread(target=update_loop, daemon=True).start()

# =============================================================================
# Gradio landing page (just a link — keeps Spaces' gradio SDK happy)
# =============================================================================
with gr.Blocks(title="Himawari PAR Nowcast") as demo:
    gr.Markdown("### 🌀 Himawari-9 Live Nowcast — Philippine Area of Responsibility")
    gr.HTML('<a href="/live" target="_self" style="display:inline-block;padding:10px 20px;'
            'background:#e63946;color:white;border-radius:6px;text-decoration:none;'
            'font-weight:bold;">🗺️ Open Live Map</a>')

demo.launch(server_name="0.0.0.0", server_port=7860, prevent_thread_lock=True)

demo.app.mount("/frames", StaticFiles(directory=FRAMES_DIR), name="frames")

# =============================================================================
# Standalone map page (real HTML response -- scripts execute normally here)
# =============================================================================
_bounds_js = json.dumps([[DISPLAY_BBOX[1],DISPLAY_BBOX[0]],[DISPLAY_BBOX[3],DISPLAY_BBOX[2]]])
_par_js = json.dumps([[lat,lon] for lon,lat in PAR_POLY])

_JS_LOGIC = """
(function(){
  const BOUNDS = __BOUNDS__;
  const PAR = __PAR__;
  const map = L.map('liveMap', {zoomControl:true});
  map.fitBounds(BOUNDS);
  L.tileLayer('https://{s}.basemaps.cartocdn.com/dark_all/{z}/{x}/{y}{r}.png', {
    attribution: '&copy; CARTO &copy; OpenStreetMap contributors',
    subdomains: 'abcd', maxZoom: 19
  }).addTo(map);
  L.rectangle(BOUNDS, {color:'#888', weight:1, dashArray:'4', fill:false}).addTo(map);
  L.polygon(PAR, {color:'#ff3b30', weight:2, fill:false}).addTo(map);

  let overlay = L.imageOverlay('/frames/obs_0.png', BOUNDS, {opacity:0.78}).addTo(map);

  let manifest = null;
  let live = true;
  let lastVersion = -1;

  const slider = document.getElementById('frameSlider');
  const label = document.getElementById('frameLabel');
  const liveBtn = document.getElementById('liveBtn');

  function frameName(idx){
    if (!manifest) return null;
    if (idx < manifest.obs_count) {
      const minsAgo = (manifest.obs_count - 1 - idx) * manifest.obs_step_min;
      return {file: 'obs_' + idx + '.png', label: 'OBS  T-' + minsAgo + 'min'};
    } else {
      const fi = idx - manifest.obs_count;
      const minsFwd = (fi + 1) * manifest.fcst_step_min;
      return {file: 'fcst_' + fi + '.png', label: 'FCST T+' + minsFwd + 'min'};
    }
  }

  function showFrame(idx){
    const fr = frameName(idx);
    if (!fr) return;
    overlay.setUrl('/frames/' + fr.file + '?t=' + Date.now());
    label.textContent = fr.label + (manifest.base_time ? ('  |  base ' + manifest.base_time) : '');
    slider.value = idx;
  }

  function applyManifest(m){
    manifest = m;
    const total = m.obs_count + m.fcst_count;
    slider.max = total - 1;
    if (live) showFrame(m.obs_count - 1);
  }

  function poll(){
    fetch('/frames/manifest.json?t=' + Date.now())
      .then(r => r.json())
      .then(m => {
        if (!m.ready) return;
        if (m.version !== lastVersion) {
          lastVersion = m.version;
          applyManifest(m);
        }
      })
      .catch(()=>{});
  }

  slider.addEventListener('input', function(){
    live = false;
    liveBtn.style.opacity = 0.45;
    showFrame(parseInt(slider.value));
  });

  liveBtn.addEventListener('click', function(){
    live = true;
    liveBtn.style.opacity = 1;
    if (manifest) showFrame(manifest.obs_count - 1);
  });

  poll();
  setInterval(poll, 20000);
})();
"""
_JS_LOGIC = _JS_LOGIC.replace("__BOUNDS__", _bounds_js).replace("__PAR__", _par_js)

FULL_PAGE_HTML = f"""<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8"/>
<title>Himawari PAR Nowcast</title>
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/leaflet/1.9.4/leaflet.css"/>
<style>body{{margin:0;font-family:monospace;background:#111;color:#eee;}}
#controls{{padding:10px;display:flex;align-items:center;gap:10px;}}</style>
</head>
<body>
<div id="liveMap" style="width:100%;height:90vh;"></div>
<div id="controls">
  <button id="liveBtn" style="padding:6px 14px;border-radius:6px;border:none;background:#e63946;color:white;font-weight:bold;cursor:pointer;">🔴 LIVE</button>
  <span id="frameLabel">Loading...</span>
  <input id="frameSlider" type="range" min="0" max="1" value="0" style="flex:1;">
</div>
<script src="https://cdnjs.cloudflare.com/ajax/libs/leaflet/1.9.4/leaflet.js"></script>
<script>
{_JS_LOGIC}
</script>
</body>
</html>"""

@demo.app.get("/live", response_class=HTMLResponse)
def live_map():
    return FULL_PAGE_HTML

demo.block_thread()