import os
import numpy as np
np.Inf = np.inf  # Patch for Keras/Numpy 2.0 compatibility
import tensorflow as tf

# Enable GPU Memory Growth to prevent TF from hoarding all 48GB of VRAM immediately
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

from data_generator import RockGlacierDataGenerator
from models.swin_unet import build_model as build_swin_unet
from models.attention_unet import build_attention_unet
from models.unet import build_unet
from models.segformer import build_segformer
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, CSVLogger, ReduceLROnPlateau

# ─────────────────────────────────────────────────────────────────────────────
# Data + band selection  (all env-driven; patch size auto-detected)
#   DATA_DIR  : patches root (default all_data/patches) -> {train,val}/{features,masks}
#   BAND_LIST : comma-separated band names, used as-is        (highest priority)
#   BANDS     : a named preset in BAND_SETS below
#   RUN_NAME  : names the checkpoint dir (all_data/checkpoint_<RUN_NAME>)
# Patch size is read from the data, so the SAME script trains the 64 px (30 m)
# set and the 128 px (10 m) set with no edits.
# ─────────────────────────────────────────────────────────────────────────────
import json as _json, glob as _glob

DATA_DIR   = os.environ.get("DATA_DIR", "all_data/patches")
TRAIN_FEAT = f"{DATA_DIR}/train/features"
TRAIN_MASK = f"{DATA_DIR}/train/masks"
VAL_FEAT   = os.environ.get("VAL_FEAT",  f"{DATA_DIR}/val/features")
VAL_MASK   = os.environ.get("VAL_MASK",  f"{DATA_DIR}/val/masks")

BAND_SETS = {
    "14": ["TPI", "ROUGHNESS", "LST_winter", "NDWI", "LST_summer", "REL_RELIEF",
           "ASPECT", "LST_anomaly", "ELEV_COP", "B2", "HILLSHADE", "B8A", "NDMI", "B12"],
    "17": ["TPI", "ROUGHNESS", "LST_winter", "NDWI", "LST_summer", "REL_RELIEF",
           "ASPECT", "LST_anomaly", "ELEV_COP", "B2", "HILLSHADE", "B8A", "NDMI", "B12",
           "B3", "B4", "B8"],
    "opt_idx": ["B2", "B3", "B4", "B8", "B8A", "B11", "B12", "NDVI", "NDWI", "NDMI", "NDSI", "BSI", "ALBEDO", "ELEV_COP"],
    "terr_therm": ["TPI", "ROUGHNESS", "SLOPE", "ASPECT", "REL_RELIEF", "HILLSHADE", "CURVATURE", "ELEV_COP", "LST_winter", "LST_summer", "LST_annual", "LST_anomaly"],
    "mixed_16": ["TPI", "ROUGHNESS", "REL_RELIEF", "ELEV_COP", "SLOPE", "LST_winter", "LST_summer", "NDWI", "NDVI", "B2", "B3", "B4", "B8", "B12", "NDMI", "BSI"],
    "top_20": ["TPI", "ROUGHNESS", "LST_winter", "NDWI", "LST_summer", "REL_RELIEF", "ASPECT", "LST_anomaly", "ELEV_COP", "B2", "HILLSHADE", "B8A", "NDMI", "B12", "B3", "B4", "B8", "NDVI", "SLOPE", "CURVATURE"],
}

_band_names = _json.load(open(f"{DATA_DIR}/metadata/norm_stats.json"))["band_names"]
_blist = os.environ.get("BAND_LIST")
_sel   = os.environ.get("BANDS", "all")
if _blist:
    _bands, RUN_NAME = [b.strip() for b in _blist.split(",") if b.strip()], os.environ.get("RUN_NAME", "custom")
elif _sel in BAND_SETS:
    _bands, RUN_NAME = BAND_SETS[_sel], os.environ.get("RUN_NAME", _sel + "band")
else:
    _bands, RUN_NAME = None, os.environ.get("RUN_NAME", "all")

if _bands is not None:
    _missing = [b for b in _bands if b not in _band_names]
    assert not _missing, f"Unknown band(s) {_missing}\nAvailable: {_band_names}"
    BAND_INDICES, N_BANDS = [_band_names.index(b) for b in _bands], len(_bands)
else:
    BAND_INDICES, N_BANDS, _bands = None, len(_band_names), list(_band_names)

CKPT_DIR = os.environ.get("CKPT_DIR",
    "all_data/checkpoint" if RUN_NAME == "all" else f"all_data/checkpoint_{RUN_NAME}")
os.makedirs(CKPT_DIR, exist_ok=True)

# Patch size auto-detected from a sample patch (64 px now, 128 px for 10 m later).
_ph, _pw = np.load(sorted(_glob.glob(f"{TRAIN_FEAT}/*.npy"))[0]).shape[:2]
print(f"[train] run={RUN_NAME}  bands={N_BANDS}  patch={_ph}x{_pw}  data={DATA_DIR}")

# ─────────────────────────────────────────────────────────────────────────────
# Hyperparameters
# ─────────────────────────────────────────────────────────────────────────────
BATCH_SIZE  = int(os.environ.get("BATCH_SIZE", 32))
EPOCHS      = int(os.environ.get("EPOCHS", 150))
LR          = 1e-4
INPUT_SHAPE = (_ph, _pw, N_BANDS)

# ─────────────────────────────────────────────────────────────────────────────
# Generators
# ─────────────────────────────────────────────────────────────────────────────
print("Initializing Data Generators...")
_COPY_PASTE = os.environ.get("COPY_PASTE", "0") == "1"
_CP_PROB    = float(os.environ.get("COPY_PASTE_PROB", 0.5))
train_gen = RockGlacierDataGenerator(TRAIN_FEAT, TRAIN_MASK, batch_size=BATCH_SIZE,
                                      target_size=INPUT_SHAPE[:2], augment=True,
                                      return_weight=True, radiometric=True,
                                      band_indices=BAND_INDICES,
                                      copy_paste=_COPY_PASTE, copy_paste_prob=_CP_PROB)
                                      # use ALL patches; COPY_PASTE=1 enables copy-paste aug
val_gen   = RockGlacierDataGenerator(VAL_FEAT, VAL_MASK, batch_size=BATCH_SIZE,
                                      target_size=INPUT_SHAPE[:2], augment=False, shuffle=False,
                                      return_weight=True, band_indices=BAND_INDICES)
                                      # val keeps ALL patches for honest evaluation

# Record the exact normalization stats used, so evaluate.py matches at test time.
import json
with open(os.path.join(CKPT_DIR, 'norm_stats.json'), 'w') as f:
    json.dump({'mean': np.asarray(train_gen.GLOBAL_MEAN).reshape(-1).tolist(),
               'std':  np.asarray(train_gen.GLOBAL_STD).reshape(-1).tolist()}, f, indent=2)

# Record the exact run config (bands, patch size) so evaluate.py reproduces it
# with zero chance of a band/order mismatch.
with open(os.path.join(CKPT_DIR, 'run_config.json'), 'w') as f:
    json.dump({'run_name': RUN_NAME, 'bands': _bands, 'n_bands': N_BANDS,
               'patch_size': [int(_ph), int(_pw)], 'data_dir': DATA_DIR}, f, indent=2)


# ─────────────────────────────────────────────────────────────────────────────
# Loss / metrics  — PADDING-MASKED.
# y_true is (B,H,W,2): channel 0 = label, channel 1 = validity weight (1=real
# pixel, 0=zero-padding). All sums below are restricted to valid pixels so the
# model is never trained or scored on padded regions.
# ─────────────────────────────────────────────────────────────────────────────
def _split_label_weight(y_true):
    yt = tf.cast(y_true[..., 0:1], tf.float32)
    w  = tf.cast(y_true[..., 1:2], tf.float32)
    return yt, w


# Focal-Tversky hyperparameters.
#   Tversky index TI = TP / (TP + ALPHA*FP + BETA*FN)
#   ALPHA>BETA penalizes false positives harder -> raises precision (the model
#   currently over-segments: FP > FN). GAMMA focuses the loss on hard pixels.
ALPHA_FP     = 0.4   # slightly higher — new dataset has real background patches in training
BETA_FN      = 0.6   # balanced — new dataset has much better class quality
GAMMA_FT     = 1.3333
BCE_WEIGHT   = 0.5     # light masked BCE term for stable early gradients
LABEL_SMOOTH = 0.05


def focal_tversky_loss(y_true, y_pred):
    smooth = 1e-6
    yt, w = _split_label_weight(y_true)
    yp = tf.cast(y_pred, tf.float32) * w          # soft probs, padded pixels zeroed
    yt = yt * w
    tp = tf.reduce_sum(yt * yp, axis=[1, 2, 3])
    fp = tf.reduce_sum((1.0 - yt) * yp, axis=[1, 2, 3])
    fn = tf.reduce_sum(yt * (1.0 - yp), axis=[1, 2, 3])
    ti = (tp + smooth) / (tp + ALPHA_FP * fp + BETA_FN * fn + smooth)
    return tf.reduce_mean(tf.pow(1.0 - ti, GAMMA_FT))


def combined_loss(y_true, y_pred):
    yt, w = _split_label_weight(y_true)
    yt_s = yt * (1.0 - LABEL_SMOOTH) + 0.5 * LABEL_SMOOTH      # binary label smoothing
    bce = tf.keras.backend.binary_crossentropy(yt_s, tf.cast(y_pred, tf.float32))
    bce = tf.reduce_sum(bce * w) / (tf.reduce_sum(w) + 1e-6)   # mean over valid pixels
    return BCE_WEIGHT * bce + focal_tversky_loss(y_true, y_pred)


# ─────────────────────────────────────────────────────────────────────────────
# Metrics  (all at threshold 0.5, valid pixels only; logged by CSVLogger)
#
# IMPORTANT: each metric uses tf.cond to return 0.0 when the batch contains
# no positive pixels, instead of the previous 1e-6/1e-6 = 1.0 artifact.
# With 52% background-only patches, the old smoothed formulas returned 1.0 for
# those batches and inflated the epoch-averaged val_dice by ~0.4–0.5 points,
# causing ModelCheckpoint to save the wrong epoch as "best".
# ─────────────────────────────────────────────────────────────────────────────
def dice_coef(y_true, y_pred):
    yt, w = _split_label_weight(y_true)
    yp = tf.cast(y_pred > 0.5, tf.float32) * w
    yt = tf.cast(yt > 0.5, tf.float32) * w
    inter = tf.reduce_sum(yt * yp)
    denom = tf.reduce_sum(yt) + tf.reduce_sum(yp)
    return tf.cond(denom < 0.5, lambda: 0.0, lambda: (2. * inter) / denom)


def precision_m(y_true, y_pred):
    yt, w = _split_label_weight(y_true)
    yp = tf.cast(y_pred > 0.5, tf.float32) * w
    yt = tf.cast(yt > 0.5, tf.float32) * w
    tp = tf.reduce_sum(yt * yp)
    fp = tf.reduce_sum(yp * (1.0 - yt))
    denom = tp + fp
    return tf.cond(denom < 0.5, lambda: 0.0, lambda: tp / denom)


def recall_m(y_true, y_pred):
    yt, w = _split_label_weight(y_true)
    yp = tf.cast(y_pred > 0.5, tf.float32) * w
    yt = tf.cast(yt > 0.5, tf.float32) * w
    tp = tf.reduce_sum(yt * yp)
    fn = tf.reduce_sum(yt * (1.0 - yp))
    denom = tp + fn
    return tf.cond(denom < 0.5, lambda: 0.0, lambda: tp / denom)


def f1_m(y_true, y_pred):
    p = precision_m(y_true, y_pred); r = recall_m(y_true, y_pred)
    denom = p + r
    return tf.cond(denom < 1e-6, lambda: 0.0, lambda: (2.0 * p * r) / denom)


def iou_m(y_true, y_pred):
    yt, w = _split_label_weight(y_true)
    yp = tf.cast(y_pred > 0.5, tf.float32) * w
    yt = tf.cast(yt > 0.5, tf.float32) * w
    inter = tf.reduce_sum(yt * yp)
    union = tf.reduce_sum(yt) + tf.reduce_sum(yp) - inter
    return tf.cond(union < 0.5, lambda: 0.0, lambda: inter / union)


def acc_m(y_true, y_pred):
    yt, w = _split_label_weight(y_true)
    yp = tf.cast(y_pred > 0.5, tf.float32)
    yt_b = tf.cast(yt > 0.5, tf.float32)
    correct = tf.cast(tf.equal(yp, yt_b), tf.float32) * w
    return tf.reduce_sum(correct) / (tf.reduce_sum(w) + 1e-6)


# ─────────────────────────────────────────────────────────────────────────────
# Build Model
# ─────────────────────────────────────────────────────────────────────────────
print("Building Model...")
MODEL_TYPE = os.environ.get("MODEL_NAME", "swin_unet")
if MODEL_TYPE == "attention_unet":
    model = build_attention_unet(input_shape=INPUT_SHAPE, num_classes=1)
elif MODEL_TYPE == "unet":
    model = build_unet(input_shape=INPUT_SHAPE, num_classes=1)
elif MODEL_TYPE == "segformer":
    model = build_segformer(input_shape=INPUT_SHAPE, num_classes=1)
else:
    model = build_swin_unet(input_shape=INPUT_SHAPE, num_classes=1, num_heads=8)
# Cosine annealing: smoothly decay LR from 1e-4 → 1e-6 over all epochs.
# Replaces ReduceLROnPlateau — no sudden drops, lets the model explore longer.
_total_steps = len(train_gen) * EPOCHS
cosine_lr = tf.keras.optimizers.schedules.CosineDecay(
    initial_learning_rate=LR,
    decay_steps=_total_steps,
    alpha=1e-2,   # minimum LR = LR * alpha = 1e-4 * 1e-2 = 1e-6
)
_WD = float(os.environ.get("WEIGHT_DECAY", 1e-4))   # raise (e.g. 5e-4) to regularize more
try:
    optimizer = tf.keras.optimizers.AdamW(learning_rate=cosine_lr, weight_decay=_WD)
    print(f"Optimizer: AdamW + CosineDecay  ({_total_steps} steps, wd={_WD})")
except Exception as e:
    optimizer = tf.keras.optimizers.Adam(learning_rate=cosine_lr)
    print(f"Optimizer: Adam + CosineDecay  ({_total_steps} steps)")
model.compile(optimizer=optimizer,
              loss=combined_loss,
              metrics=[acc_m, dice_coef, precision_m, recall_m, f1_m, iou_m])

model.summary()

# ─────────────────────────────────────────────────────────────────────────────
# Callbacks
# ─────────────────────────────────────────────────────────────────────────────
callbacks = [
    ModelCheckpoint(os.path.join(CKPT_DIR, 'best_model.weights.h5'),
                    monitor='val_loss', mode='min',
                    save_best_only=True, save_weights_only=True, verbose=1),
    # Patience=35: cosine LR decays smoothly so the model needs more room to
    # naturally improve without the crutch of sudden LR reductions.
    EarlyStopping(monitor='val_loss', patience=35, restore_best_weights=True, verbose=1),
    # ReduceLROnPlateau removed — cosine annealing handles LR schedule.
    CSVLogger(os.path.join(CKPT_DIR, 'training_history.csv')),
]

# ─────────────────────────────────────────────────────────────────────────────
# Train
# ─────────────────────────────────────────────────────────────────────────────
print("Starting Training...")
history = model.fit(
    train_gen,
    validation_data=val_gen,
    epochs=EPOCHS,
    callbacks=callbacks,
)
print("Training Complete.")