"""
校准曲线可视化：读取 output/02_分组校准数据 下的 CSV，生成交互式 HTML。

用法:
    python tasks/承接/rosn校准/plot_calibration.py
"""

import os
import json
import glob
import pandas as pd
from collections import defaultdict

DATA_DIRS = [
    os.path.join(os.path.dirname(__file__), "output", "02_分组校准数据"),
    os.path.join(os.path.dirname(__file__), "output", "03_label分桶校准数据"),
]
OUT_HTML = os.path.join(os.path.dirname(__file__), "output", "calibration.html")


def load_data():
    files = []
    for d in DATA_DIRS:
        files.extend(sorted(glob.glob(os.path.join(d, "*.csv"))))
    if not files:
        raise FileNotFoundError(f"No CSV files found in {DATA_DIRS}")
    df = pd.concat([pd.read_csv(f) for f in files], ignore_index=True)
    df["dt"] = df["dt"].astype(str)
    return df


def build_charts_data(df):
    """
    扁平化数据结构:
    ALL_DATA[apptype][bucket_type][group_key][dt] = [ {bucket_id, predict, stat, label, ...}, ... ]
    """
    all_data = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(list))))
    for _, row in df.iterrows():
        at = str(row["apptype"])
        bt = row["bucket_type"]
        gk = row["group_key"]
        dt = row["dt"]
        predict = float(row["predict"])
        all_data[at][bt][gk][dt].append({
            "bucket_id": int(row["bucket_id"]),
            "predict": predict,
            "pred_校准": round(1.22 * predict ** 1.15, 6),
            "stat": float(row["stat"]) if pd.notna(row.get("stat")) else None,
            "label": float(row["label"]),
            "cnt": int(row["cnt"]),
            "range_begin": float(row["range_begin"]),
            "range_end": float(row["range_end"]),
        })
    # 排序 rows by bucket_id
    for at in all_data:
        for bt in all_data[at]:
            for gk in all_data[at][bt]:
                for dt in all_data[at][bt][gk]:
                    all_data[at][bt][gk][dt].sort(key=lambda r: r["bucket_id"])
    return all_data


def render_html(all_data, all_dts):
    """生成单图 ECharts HTML — 动态配色，5 项筛选"""

    # 收集维度值
    all_apptypes = sorted(all_data.keys())
    all_bucket_types = sorted({bt for at in all_data for bt in all_data[at]})
    all_group_keys = sorted({gk for at in all_data for bt in all_data[at] for gk in all_data[at][bt]})

    # apptype 标签映射
    apptype_labels = {}
    for at in all_apptypes:
        apptype_labels[at] = "视频号" if at == "4" else f"apptype={at}"

    # ---- 序列化数据 ----
    js_all_data = {}
    for at in all_data:
        js_all_data[at] = {}
        for bt in all_data[at]:
            js_all_data[at][bt] = {}
            for gk in all_data[at][bt]:
                js_all_data[at][bt][gk] = {}
                for dt, rows in all_data[at][bt][gk].items():
                    js_all_data[at][bt][gk][dt] = rows

    # ---- 构建筛选控件 HTML ----
    dt_options = ''.join(
        f'<option value="{dt}">{dt}</option>' for dt in sorted(all_dts, reverse=True)
    )
    gk_checkboxes = ''.join(
        '<label class="cb-item">'
        f'<input type="checkbox" value="{gk}" {"checked" if gk == "对照组" else ""} onchange="renderChart()">'
        f'{gk}</label>'
        for gk in all_group_keys
    )
    bt_checkboxes = ''.join(
        f'<label class="cb-item"><input type="checkbox" value="{bt}" checked onchange="renderChart()">{bt}</label>'
        for bt in all_bucket_types
    )
    at_checkboxes = ''.join(
        '<label class="cb-item">'
        f'<input type="checkbox" value="{at}" checked onchange="renderChart()">'
        f'{apptype_labels[at]}</label>'
        for at in all_apptypes
    )

    html = f"""<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<title>校准曲线</title>
<script src="https://cdn.jsdelivr.net/npm/echarts@5/dist/echarts.min.js"></script>
<style>
    body {{ font-family: -apple-system, sans-serif; margin: 20px; background: #fafafa; }}
    .filter-row {{ display: flex; align-items: center; gap: 12px; margin-bottom: 12px; }}
    .filter-row > label {{ font-weight: 600; min-width: 56px; }}
    select {{ padding: 4px 8px; border-radius: 4px; border: 1px solid #ccc; }}
    .checkbox-group {{ display: flex; flex-wrap: wrap; gap: 6px 14px; }}
    .cb-item {{ display: inline-flex; align-items: center; gap: 4px; cursor: pointer; font-size: 13px; }}
    .cb-line {{ display: inline-block; width: 18px; height: 0; margin-bottom: 1px; }}
    #chart-container {{ background: #fff; border-radius: 8px; padding: 16px; box-shadow: 0 1px 3px rgba(0,0,0,.1); margin-top: 16px; }}
    #mae-info {{ background: #fff; border-radius: 8px; padding: 16px; box-shadow: 0 1px 3px rgba(0,0,0,.1); margin-top: 12px; font-size: 13px; line-height: 1.6; }}
    #mae-info table {{ font-size: 12px; font-variant-numeric: tabular-nums; }}
    #mae-info th {{ font-weight: 600; white-space: nowrap; }}
    #mae-info td {{ white-space: nowrap; }}
</style>
</head>
<body>
<h2>校准曲线</h2>

<div class="filter-row">
    <label>dt:</label>
    <select id="sel_dt" onchange="renderChart()">{dt_options}</select>
</div>
<div class="filter-row">
    <label>分组:</label>
    <span id="gk_checks" class="checkbox-group">{gk_checkboxes}</span>
</div>
<div class="filter-row">
    <label>指标:</label>
    <span id="metric_checks" class="checkbox-group">
        <label class="cb-item"><input type="checkbox" value="predict" checked onchange="renderChart()"><span class="cb-line" style="border-top:2px solid #333"></span>predict</label>
        <label class="cb-item"><input type="checkbox" value="pred_校准" checked onchange="renderChart()"><span class="cb-line" style="border-top:2px dashed #333"></span>pred_校准</label>
        <label class="cb-item"><input type="checkbox" value="stat" checked onchange="renderChart()"><span class="cb-line" style="border-top:2px dashed #333"></span>stat</label>
        <label class="cb-item"><input type="checkbox" value="label" checked onchange="renderChart()"><span class="cb-line" style="border-top:2px dotted #333"></span>label</label>
    </span>
</div>
<div class="filter-row">
    <label>分桶:</label>
    <span id="bt_checks" class="checkbox-group">{bt_checkboxes}</span>
</div>
<div class="filter-row">
    <label>apptype:</label>
    <span id="at_checks" class="checkbox-group">{at_checkboxes}</span>
</div>

<div id="chart-container">
    <div id="main-chart" style="width:100%;height:600px;"></div>
</div>
<div id="mae-info"></div>

<script>
var ALL_DATA = {json.dumps(js_all_data, ensure_ascii=False)};
var APPTYPE_LABELS = {json.dumps(apptype_labels, ensure_ascii=False)};

var PALETTE = [
    '#3b82f6', '#f97316', '#22c55e', '#ef4444', '#a855f7',
    '#06b6d4', '#f59e0b', '#ec4899', '#64748b', '#84cc16',
    '#0ea5e9', '#d946ef', '#14b8a6', '#f43f5e', '#8b5cf6',
    '#e11d48', '#059669', '#7c3aed', '#ea580c', '#0284c7',
];

var LINE_STYLES = {{
    'predict': 'solid',
    'pred_校准': [10, 3],
    'stat':    [6, 3],
    'label':   [2, 2]
}};

var chart = echarts.init(document.getElementById('main-chart'));

function getCheckedValues(id) {{
    var vals = [];
    document.getElementById(id).querySelectorAll('input[type=checkbox]').forEach(function(cb) {{
        if (cb.checked) vals.push(cb.value);
    }});
    return vals;
}}

function renderChart() {{
    var dt = document.getElementById('sel_dt').value;
    var selGroups = getCheckedValues('gk_checks');
    var selMetrics = getCheckedValues('metric_checks');
    var selBucketTypes = getCheckedValues('bt_checks');
    var selApptypes = getCheckedValues('at_checks');

    // 1. 构建 combo 列表 (apptype × bucket_type × group_key)
    var combos = [];
    selApptypes.forEach(function(at) {{
        selBucketTypes.forEach(function(bt) {{
            selGroups.forEach(function(gk) {{
                if (ALL_DATA[at] && ALL_DATA[at][bt] && ALL_DATA[at][bt][gk]) {{
                    combos.push({{ at: at, bt: bt, gk: gk }});
                }}
            }});
        }});
    }});

    // 2. 动态配色: 每个 combo 分配一个颜色
    var series = [];
    var bucketIds = null;

    combos.forEach(function(combo, idx) {{
        var color = PALETTE[idx % PALETTE.length];
        var rows = (ALL_DATA[combo.at][combo.bt][combo.gk][dt]) || [];
        if (!bucketIds && rows.length > 0) {{
            bucketIds = rows.map(function(r) {{ return r.bucket_id; }});
        }}
        var atLabel = APPTYPE_LABELS[combo.at] || combo.at;
        var comboLabel = atLabel + '·' + combo.bt + '·' + combo.gk;

        // 3. 每个 combo × selMetrics → 一条 series
        selMetrics.forEach(function(metric) {{
            series.push({{
                name: comboLabel + '·' + metric,
                type: 'line',
                data: rows.map(function(r) {{ return r[metric]; }}),
                symbol: 'none',
                lineStyle: {{
                    type: LINE_STYLES[metric],
                    color: color,
                    width: metric === 'label' ? 2 : 1.5
                }},
                itemStyle: {{ color: color }}
            }});
        }});
    }});

    chart.setOption({{
        tooltip: {{
            trigger: 'axis',
            formatter: function(params) {{
                if (!params.length) return '';
                var dataIdx = params[0].dataIndex;
                var s = 'bucket: ' + (bucketIds ? bucketIds[dataIdx] : dataIdx) + '<br/>';
                // 按 combo 分组显示，从 ALL_DATA 取原始行
                var groups = {{}};
                params.forEach(function(p) {{
                    var parts = p.seriesName.split('·');
                    var metric = parts.pop();
                    var comboKey = parts.join('·');
                    if (!groups[comboKey]) groups[comboKey] = {{}};
                    groups[comboKey][metric] = p.value;
                    groups[comboKey].color = p.color;
                }});
                combos.forEach(function(combo, ci) {{
                    var atLabel = APPTYPE_LABELS[combo.at] || combo.at;
                    var ck = atLabel + '·' + combo.bt + '·' + combo.gk;
                    var g = groups[ck];
                    if (!g) return;
                    var rows = (ALL_DATA[combo.at][combo.bt][combo.gk][dt]) || [];
                    var r = rows[dataIdx];
                    s += '<br/><b style="color:' + g.color + '">' + ck + '</b>';
                    if (r) s += ' &nbsp;<span style="color:#999">cnt=' + r.cnt + ' range=[' + r.range_begin.toFixed(4) + ', ' + r.range_end.toFixed(4) + ']</span>';
                    s += '<br/>';
                    if (g.predict != null) s += '  predict: ' + g.predict.toFixed(4) + '<br/>';
                    if (g['pred_校准'] != null) s += '  pred_校准: ' + g['pred_校准'].toFixed(4) + '<br/>';
                    if (g.stat != null) s += '  stat: ' + g.stat.toFixed(4) + '<br/>';
                    if (g.label != null) s += '  label: ' + g.label.toFixed(4) + '<br/>';
                }});
                return s;
            }}
        }},
        legend: {{ type: 'scroll', bottom: 0 }},
        grid: {{ left: 60, right: 30, top: 30, bottom: 60 }},
        xAxis: {{ name: 'bucket_id', type: 'category', data: bucketIds || [], nameLocation: 'center', nameGap: 28 }},
        yAxis: {{ name: 'value', type: 'value', nameLocation: 'center', nameGap: 45 }},
        series: series
    }}, true);

    // 计算并展示 MAE（label 加权，总体 + 分段）
    function calcWMAE(rows, metric) {{
        var wSum = 0, wTotal = 0;
        rows.forEach(function(r) {{
            if (r[metric] != null && r.label != null) {{
                wSum += r.label * Math.abs(r[metric] - r.label);
                wTotal += r.label;
            }}
        }});
        return wTotal > 0 ? (wSum / wTotal) : null;
    }}
    function rangeInfo(rows) {{
        var labelMin = Infinity, labelMax = -Infinity, rngMin = Infinity, rngMax = -Infinity;
        rows.forEach(function(r) {{
            if (r.label != null) {{ labelMin = Math.min(labelMin, r.label); labelMax = Math.max(labelMax, r.label); }}
            if (r.range_begin != null) rngMin = Math.min(rngMin, r.range_begin);
            if (r.range_end != null) rngMax = Math.max(rngMax, r.range_end);
        }});
        return 'label=[' + labelMin.toFixed(4) + ', ' + labelMax.toFixed(4) + ']'
            + ' bucket=[' + rngMin.toFixed(4) + ', ' + rngMax.toFixed(4) + ']';
    }}
    function maeRow(metrics, rows) {{
        var maeMap = {{}};
        metrics.forEach(function(m) {{
            maeMap[m] = calcWMAE(rows, m);
        }});
        // 校准Δ
        var delta = null, deltaPct = null;
        if (maeMap['predict'] != null && maeMap['pred_校准'] != null) {{
            delta = maeMap['pred_校准'] - maeMap['predict'];
            deltaPct = delta / maeMap['predict'] * 100;
        }}
        return {{ maeMap: maeMap, delta: delta, deltaPct: deltaPct, rng: rangeInfo(rows) }};
    }}
    function fmtVal(v) {{ return v != null ? v.toFixed(4) : '-'; }}
    function fmtDelta(d) {{
        if (d == null) return '-';
        var sign = d.delta <= 0 ? '' : '+';
        var color = d.delta <= 0 ? '#16a34a' : '#dc2626';
        return '<span style="color:' + color + '">' + sign + d.delta.toFixed(4) + ' (' + sign + d.deltaPct.toFixed(1) + '%)</span>';
    }}

    var maeHtml = '';
    var maeMetrics = selMetrics.filter(function(m) {{ return m !== 'label'; }});
    if (maeMetrics.length > 0) {{
        // 表头
        var thMetrics = '';
        maeMetrics.forEach(function(m) {{ thMetrics += '<th>' + m + '</th>'; }});
        var hasCalDelta = maeMetrics.indexOf('predict') >= 0 && maeMetrics.indexOf('pred_校准') >= 0;

        combos.forEach(function(combo, idx) {{
            var color = PALETTE[idx % PALETTE.length];
            var rows = (ALL_DATA[combo.at][combo.bt][combo.gk][dt]) || [];
            if (rows.length === 0) return;
            var atLabel = APPTYPE_LABELS[combo.at] || combo.at;
            var comboLabel = atLabel + '·' + combo.bt + '·' + combo.gk;

            var n = rows.length;
            var cut1 = Math.floor(n / 3), cut2 = Math.floor(n * 2 / 3);
            var segments = [
                {{ name: '总体', rows: rows }},
                {{ name: '低(0~' + (cut1 - 1) + ')', rows: rows.slice(0, cut1) }},
                {{ name: '中(' + cut1 + '~' + (cut2 - 1) + ')', rows: rows.slice(cut1, cut2) }},
                {{ name: '高(' + cut2 + '~' + (n - 1) + ')', rows: rows.slice(cut2) }}
            ];

            maeHtml += '<div style="margin-bottom:12px"><b style="color:' + color + '">' + comboLabel + '</b>';
            maeHtml += '<table style="border-collapse:collapse;margin-top:4px;width:100%"><tr style="background:#f5f5f5">';
            maeHtml += '<th style="text-align:left;padding:3px 8px">分段</th>' + thMetrics.replace(/<th>/g, '<th style="padding:3px 8px">');
            if (hasCalDelta) maeHtml += '<th style="padding:3px 8px">校准Δ</th>';
            maeHtml += '<th style="padding:3px 8px">label范围</th><th style="padding:3px 8px">bucket范围</th></tr>';

            segments.forEach(function(seg) {{
                var d = maeRow(maeMetrics, seg.rows);
                maeHtml += '<tr style="border-top:1px solid #eee"><td style="padding:3px 8px">' + seg.name + '</td>';
                maeMetrics.forEach(function(m) {{
                    maeHtml += '<td style="padding:3px 8px;text-align:right">' + fmtVal(d.maeMap[m]) + '</td>';
                }});
                if (hasCalDelta) maeHtml += '<td style="padding:3px 8px;text-align:right">' + fmtDelta(d) + '</td>';
                maeHtml += '<td style="padding:3px 8px;color:#888">' + d.rng.split(' ')[0] + '</td>';
                maeHtml += '<td style="padding:3px 8px;color:#888">' + d.rng.split(' ')[1] + '</td>';
                maeHtml += '</tr>';
            }});
            maeHtml += '</table></div>';
        }});
    }}
    document.getElementById('mae-info').innerHTML = maeHtml;
}}

renderChart();
window.addEventListener('resize', function() {{ chart.resize(); }});
</script>
</body>
</html>"""

    os.makedirs(os.path.dirname(OUT_HTML), exist_ok=True)
    with open(OUT_HTML, "w", encoding="utf-8") as f:
        f.write(html)
    print(f"HTML saved to: {OUT_HTML}")


def main():
    df = load_data()
    all_data = build_charts_data(df)
    all_dts = sorted(df["dt"].unique())
    render_html(all_data, all_dts)


if __name__ == "__main__":
    main()
