#!/usr/bin/env python
# predict.py
import sys
import pickle
import numpy as np
from train import build_features

NON_CANCER = [0, 2]
CANCER     = [3, 5, 6]
BCKG       = [8]

# Boost dcis and indc raw probabilities before argmax.
# These multipliers shift the decision boundary without retraining.
# dcis: aggressively boost since recall is only 41%
# indc: moderate boost since it's also underperforming
SCORE_MULTIPLIERS = {
    5: 11.5,   # dcis
    6: 4.5,   # indc
    0: 1.3,   # norm
    3: 8.0,    #infl
}

def load_csv(filepath):
    labels, features = [], []
    with open(filepath, 'r') as f:
        lines = f.readlines()
    for line in lines[1:]:
        parts = line.strip().split(',')
        labels.append(int(float(parts[0])))
        features.append(np.array([float(x) for x in parts[1:]], dtype=np.float32))
    return np.array(labels), np.array(features, dtype=np.float32)


def predict(models, Xf):
    m_s1   = models['stage1']
    inv_s1 = models['s1_inv']
    m_nc   = models['nc']
    inv_nc = models['nc_inv']
    m_c    = models['cancer']
    inv_c  = models['c_inv']
    m_flat    = models['flat']
    inv_flat  = models['flat_inv']

    # ── Hierarchical probs ────────────────────────────────────────────────
    p_s1 = m_s1.predict_proba(Xf)
    s1_nc_idx = next(k for k, v in inv_s1.items() if v == 0)
    s1_c_idx  = next(k for k, v in inv_s1.items() if v == 1)
    s1_bg_idx = next(k for k, v in inv_s1.items() if v == 2)

    p_nc = m_nc.predict_proba(Xf)
    p_c  = m_c.predict_proba(Xf)

    hier_probs = np.zeros((len(Xf), 9), dtype=np.float32)
    for j, cls in inv_nc.items():
        hier_probs[:, cls] += p_nc[:, j] * p_s1[:, s1_nc_idx]
    for j, cls in inv_c.items():
        hier_probs[:, cls] += p_c[:, j] * p_s1[:, s1_c_idx]
    hier_probs[:, 8] += p_s1[:, s1_bg_idx]

    # ── Flat model probs ──────────────────────────────────────────────────
    p_flat_raw = m_flat.predict_proba(Xf)
    flat_probs = np.zeros((len(Xf), 9), dtype=np.float32)
    for j, cls in inv_flat.items():
        flat_probs[:, cls] = p_flat_raw[:, j]

    # ── Blend flat into hierarchical for hard classes only ────────────────
    BLEND = {
        5: 0.5,   # dcis  — flat model is heavily dcis-biased, use it 50%
        3: 0.3,   # infl
        6: 0.3,   # indc
    }
    final_probs = hier_probs.copy()
    for cls, w in BLEND.items():
        final_probs[:, cls] = (1 - w) * hier_probs[:, cls] + w * flat_probs[:, cls]

    # ── Score multipliers ─────────────────────────────────────────────────
    for cls, mult in SCORE_MULTIPLIERS.items():
        final_probs[:, cls] *= mult

    return np.argmax(final_probs, axis=1).astype(int)


def main():
    model_file, input_csv, output_csv = sys.argv[1:]

    print("Loading model...")
    with open(model_file, 'rb') as f:
        models = pickle.load(f)

    print("Loading data...")
    labels, X = load_csv(input_csv)

    print("Building features...")
    Xf = build_features(X)

    print("Predicting...")
    preds = predict(models, Xf)

    assert len(preds) == len(labels), f"Length mismatch: {len(preds)} vs {len(labels)}"

    print("Writing output...")
    with open(output_csv, 'w') as f:
        f.write("label\n")
        for p in preds:
            f.write(f"{p}\n")

    print("Done.")

if __name__ == "__main__":
    main()