#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
DepthART 两路深度输出的相关性/一致性对比。

支持三种输入格式，自动识别：
    - 汇总 NPZ（names/depths），如 ONNX 脚本 --dump-npz 生成的文件
    - 分块汇总目录（*.npz 中含 names/depths）
    - 旧的逐图 NPZ 目录（<图片名>.npz 中含 depth）

指标（逐图计算后汇总 mean/p50/p90/min）：
    pearson      皮尔逊相关系数
    spearman     斯皮尔曼秩相关
    cosine       余弦相似度
    mae / rmse   原始量纲误差
    aligned_*    最小二乘 scale-shift 对齐后的误差
    rel_error    mae / mean(|ref|)

用法：
    ../../venv/bin/python src/tool/compare_depth_outputs.py \
        --ref results/nyu654/fp32/npz --pred results/nyu654/axmodel \
        --out results/nyu654/eval/corr_fp32_axmodel.json
"""

import argparse
import json
from pathlib import Path

import numpy as np
from scipy.stats import spearmanr
from tqdm import tqdm


def load_predictions(path):
    path = Path(path)
    if path.is_file():
        files = [path]
    elif path.is_dir():
        files = sorted(path.glob('*.npz'))
        if not files:
            raise ValueError(f'目录中没有 NPZ: {path}')
    else:
        raise FileNotFoundError(f'路径不存在: {path}')

    preds = {}
    for fp in files:
        with np.load(fp, allow_pickle=False) as data:
            if 'depths' in data and 'names' in data:
                for name, depth in zip(data['names'], data['depths']):
                    preds[str(name)] = np.asarray(depth)
            elif 'depth' in data:
                preds[fp.name[:-len('.npz')]] = np.asarray(data['depth'])
            else:
                raise ValueError(f'无法识别的 NPZ: {fp}')
    return preds


def compare(ref, pred):
    a = np.asarray(ref, dtype=np.float64).ravel()
    b = np.asarray(pred, dtype=np.float64).ravel()
    if a.shape != b.shape:
        raise ValueError(f'shape 不一致: {a.shape} vs {b.shape}')
    diff = b - a
    mae = float(np.mean(np.abs(diff)))
    rmse = float(np.sqrt(np.mean(diff ** 2)))
    cosine = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-12))
    pearson = float(np.corrcoef(a, b)[0, 1])
    spearman = float(spearmanr(a, b).statistic)
    slope, offset = np.polyfit(a, b, 1)
    aligned = b - (slope * a + offset)
    return {
        'pearson': pearson,
        'spearman': spearman,
        'cosine': cosine,
        'mae': mae,
        'rmse': rmse,
        'rel_error': mae / (np.mean(np.abs(a)) + 1e-12),
        'scale': float(slope),
        'offset': float(offset),
        'aligned_mae': float(np.mean(np.abs(aligned))),
        'aligned_rmse': float(np.sqrt(np.mean(aligned ** 2))),
    }


def summarize(rows):
    keys = ('pearson', 'spearman', 'cosine', 'rel_error', 'aligned_rmse')
    out = {}
    for k in keys:
        values = np.array([r[k] for r in rows])
        out[k] = {'mean': float(values.mean()), 'p50': float(np.percentile(values, 50)),
                  'p90': float(np.percentile(values, 90)), 'min': float(values.min())}
    return out


def main():
    parser = argparse.ArgumentParser(description='两路深度输出相关性对比')
    parser.add_argument('--ref', required=True, help='参考输出（基准后端）')
    parser.add_argument('--pred', required=True, help='待比较输出')
    parser.add_argument('--out', help='结果 JSON 路径')
    args = parser.parse_args()

    ref = load_predictions(args.ref)
    pred = load_predictions(args.pred)
    names = sorted(set(ref) & set(pred))
    if not names:
        raise ValueError('两路输出没有同名样本')
    missing = (len(ref) - len(names)) + (len(pred) - len(names))
    if missing:
        print(f'警告: {missing} 个样本名称不匹配，已忽略')

    rows = []
    for name in tqdm(names, desc='compare', ncols=80):
        rows.append({'name': name, **compare(ref[name], pred[name])})

    summary = summarize(rows)
    low = [r for r in rows if r['pearson'] < 0.9]
    worst = sorted(rows, key=lambda r: r['pearson'])[:5]
    report = {
        'ref': str(args.ref), 'pred': str(args.pred),
        'num_samples': len(rows),
        'summary': summary,
        'pearson_below_0.9': len(low),
        'worst': [{'name': r['name'], 'pearson': r['pearson'],
                   'mae': r['mae'], 'aligned_rmse': r['aligned_rmse']} for r in worst],
        'per_image': rows,
    }
    print(f'\n样本数 {len(rows)}，pearson<0.9 的 {len(low)} 张')
    print(f'{"指标":<14}{"mean":>12}{"p50":>12}{"p90":>12}{"min":>12}')
    for k, v in summary.items():
        print(f'{k:<14}{v["mean"]:>12.4f}{v["p50"]:>12.4f}{v["p90"]:>12.4f}{v["min"]:>12.4f}')
    print('\n最差 5 张:')
    for r in worst:
        print(f'  {r["name"]:<32} pearson={r["pearson"]:.4f} mae={r["mae"]:.4f}')

    if args.out:
        Path(args.out).parent.mkdir(parents=True, exist_ok=True)
        with open(args.out, 'w') as fp:
            json.dump(report, fp, indent=2)
        print(f'\n结果已保存: {args.out}')


if __name__ == '__main__':
    main()