#!/usr/bin/env python3
import sys
import json
import argparse
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]

def calculate_distribution_risk(row: dict, kospi_ret_5d: float) -> dict:
    close = float(row.get("Close") or row.get("close") or 0.0)
    ma20 = float(row.get("MA20") or row.get("ma20") or 0.0)
    high = float(row.get("High") or row.get("high") or close)
    low = float(row.get("Low") or row.get("low") or close)
    volume = row.get("Volume") or row.get("volume")
    avg_vol_5d = row.get("AvgVolume5d") or row.get("avg_volume_5d") or row.get("Avg_Volume_5D")
    flow_credit = row.get("FlowCredit") or row.get("flow_credit") or row.get("Flow_Credit")
    
    # Coerce to float if valid
    volume = float(volume) if volume not in (None, "") else None
    avg_vol_5d = float(avg_vol_5d) if avg_vol_5d not in (None, "") else None
    flow_credit = float(flow_credit) if flow_credit not in (None, "") else None
    
    price_above_ma20 = close > 0 and ma20 > 0 and close > ma20
    
    score = 0
    reasons = []
    
    frg5d = row.get("Frg5D") or row.get("frg_5d") or row.get("Frg_5D")
    inst5d = row.get("Inst5D") or row.get("inst_5d") or row.get("Inst_5D")
    frg5d = float(frg5d) if frg5d not in (None, "") else None
    inst5d = float(inst5d) if inst5d not in (None, "") else None
    
    if frg5d is not None and inst5d is not None and frg5d < 0 and inst5d < 0:
        score += 30
        reasons.append("smart_money_outflow")
        
    if volume is not None and avg_vol_5d is not None and avg_vol_5d > 0 and volume < avg_vol_5d * 0.80:
        score += 20
        reasons.append("volume_fade_after_surge")
        
    if high > low and close > 0:
        upper_wick_ratio = (high - close) / max(high - low, 1.0)
        if upper_wick_ratio >= 0.45 and price_above_ma20:
            score += 15
            reasons.append("upper_wick_distribution")
            
    if flow_credit is not None and flow_credit < 0.40:
        score += 20
        reasons.append("flow_credit_low")
        
    ret5d = row.get("Ret5D") or row.get("ret_5d") or row.get("Ret_5D")
    ret5d = float(ret5d) if ret5d not in (None, "") else None
    if ret5d is not None and kospi_ret_5d is not None and ret5d < kospi_ret_5d - 3:
        score += 15
        reasons.append("sector_relative_lag")
        
    ac_gate = row.get("acGate") or row.get("ac_gate")
    if ac_gate and "CLIMAX" in str(ac_gate).upper():
        score += 15
        reasons.append("anti_climax_gate")
        
    ac_total = row.get("acTotal") or row.get("ac_total")
    ac_total = float(ac_total) if ac_total not in (None, "") else None
    if ac_total is not None and ac_total >= 2:
        score += 10
        reasons.append("ac_total_gte2")
        
    val_surge_pct = row.get("valSurgePct") or row.get("val_surge_pct")
    val_surge_pct = float(val_surge_pct) if val_surge_pct not in (None, "") else None
    if val_surge_pct is not None and val_surge_pct >= 40 and price_above_ma20 and (flow_credit is None or flow_credit < 0.50):
        score += 10
        reasons.append("val_surge_no_flow_support")
        
    high52w = row.get("high52w") or row.get("high_52w") or row.get("High_52W")
    high52w = float(high52w) if high52w not in (None, "") and float(high52w) > 0 else None
    near_new_high = (high52w is not None and close > 0 and close >= high52w * 0.97) or (ma20 > 0 and close > ma20 * 1.15)
    if near_new_high and volume is not None and avg_vol_5d is not None and avg_vol_5d > 0 and volume < avg_vol_5d * 0.80:
        score += 12
        reasons.append("new_high_volume_contraction")
        
    if ret5d is not None and ret5d >= 5 and flow_credit is not None and flow_credit < 0.45:
        score += 10
        reasons.append("surge_weak_flow")
        
    final_score = min(100, max(0, score))
    state = "BLOCK_BUY" if final_score >= 70 else "TRIM_REVIEW" if final_score >= 55 else "PASS"
    pre_dist_warning = "EARLY_WARNING" if ("new_high_volume_contraction" in reasons or "surge_weak_flow" in reasons) else "NONE"
    
    return {
        "distribution_risk_score": final_score,
        "anti_distribution_state": state,
        "pre_distribution_warning": pre_dist_warning,
        "reason_codes": reasons
    }

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--json", default="GatherTradingData.json")
    parser.add_argument("--out", default="Temp/distribution_risk_score_v2.json")
    args = parser.parse_args()

    json_path = ROOT / args.json
    if not json_path.exists():
        print(f"Input file not found: {json_path}")
        sys.exit(1)

    raw = json.loads(json_path.read_text(encoding="utf-8"))
    data_feed = raw.get("data", {}).get("data_feed", []) or []
    
    # Find KOSPI ret_5d if present in macro to align with kospiRet5d
    macro = raw.get("data", {}).get("macro", []) or []
    kospi_ret_5d = 0.0
    for m in macro:
        if str(m.get("Ticker") or m.get("ticker")).strip() == "KOSPI":
            kospi_ret_5d = float(m.get("Ret5D") or m.get("ret_5d") or 0.0)
            break

    scores = {}
    for row in data_feed:
        ticker = row.get("Ticker") or row.get("ticker")
        if not ticker:
            continue
        res = calculate_distribution_risk(row, kospi_ret_5d)
        scores[ticker] = res

    out_path = ROOT / args.out
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps({
        "formula_id": "DISTRIBUTION_RISK_SCORE_V2",
        "scores": scores
    }, indent=2, ensure_ascii=False), encoding="utf-8")
    
    print(f"Saved distribution risk scores to {out_path}")
    sys.exit(0)

if __name__ == "__main__":
    main()