import os
import time
from collections import Counter, defaultdict

import numpy as np
from rich import box
from rich.columns import Columns
from rich.panel import Panel
from rich.table import Table
from rich.console import Console, Group, Text
import onnxruntime as ort
import onnx_graphsurgeon as gs

from convdog.core.graph import ConvDogModel
from convdog.utils.logger import logger



class ModelStats:
    """保存并提取模型统计信息的结构体"""

    def __init__(self, model: ConvDogModel, name: str):
        self.name = name
        # 假设你的 ConvDogGraph 现在内部维护了一个 gs_graph 属性
        # 如果没有，可以通过 gs.import_onnx(graph_wrapper.model) 转换
        g = model.graph
        self.graph = g
        self.model = model

        self.opset = g.opset
        self.ir_version = model.model.ir_version
        self.size_mb = model.model.ByteSize() / (1024 * 1024)

        # 1. 统计算子数量 (GS 中 node.op 代替了 node.op_type)
        self.op_counts = Counter([node.op for node in g.nodes])

        # 2. 提取所有的 Initializers (GS 中叫 Constants)
        # gs.Constant 对象包含 .values (numpy数组) 和 .name
        all_constants = [t for t in g.tensors().values() if isinstance(t, gs.Constant)]

        # 3. 统计权重精度与对齐性 (TensorCore Alignment)
        self.weight_counts = Counter([str(c.values.dtype) for c in all_constants])

        self.aligned_count = 0
        self.total_weight_layers = len(all_constants)
        self.total_params = 0
        self.dtype_params = defaultdict(int)

        for c in all_constants:
            arr = c.values
            # 统计总参数量
            self.total_params += arr.size
            self.dtype_params[str(arr.dtype)] += arr.size

            # 统计对齐性 (dims % 8 == 0)
            if all(d % 8 == 0 for d in arr.shape if d > 1):
                self.aligned_count += 1

        # 4. 统计输入输出 (GS 中直接是 graph.inputs / graph.outputs)
        self.inputs = self._parse_gs_io(g.inputs)
        self.outputs = self._parse_gs_io(g.outputs)

    @staticmethod
    def _parse_gs_io(io_list):
        """解析 GS 的输入输出变量"""
        info = []
        for x in io_list:
            # x 在这里是一个 gs.Variable 对象
            shape_str = []
            for d in x.shape:
                if isinstance(d, int):
                    shape_str.append(str(d))
                elif isinstance(d, str):
                    shape_str.append(d) # 符号维度
                else:
                    shape_str.append("?")
            info.append(f"{x.name}: ({', '.join(shape_str)})")
        return info


def get_diff_str(old_val, new_val, is_size=False):
    """根据数值变化生成带箭头的字符串"""
    if new_val == old_val:
        return f"{new_val:.2f} MB" if is_size else str(new_val)

    increased = new_val > old_val
    if is_size:
        color = "red" if increased else "green"
        arrow = "↑" if increased else "↓"
        return f"{new_val:.2f} MB [bold {color}]{arrow}[/]"
    else:
        color = "yellow" if increased else "green"
        arrow = "↑" if increased else "↓"
        return f"{new_val} [bold {color}]{arrow}[/]"


def generate_random_inputs(sess):
    try:
        input_dict = {}
        # 1. 动态获取 Session 需要的输入
        for input_meta in sess.get_inputs():
            name = input_meta.name
            shape = input_meta.shape
            dtype_str = input_meta.type  # e.g. 'tensor(float16)'

            # 处理动态 Shape (None 或 str)
            fixed_shape = []
            for s in shape:
                if isinstance(s, int) and s > 0:
                    fixed_shape.append(s)
                else:
                    # 如果用户在 O0 没指定 shapes，这里兜底用 1
                    fixed_shape.append(1)

            # 类型映射
            if "float16" in dtype_str:
                target_dtype = np.float16
            elif "int64" in dtype_str:
                target_dtype = np.int64
            else:
                target_dtype = np.float32

            # 生成随机数据
            data = np.random.randn(*fixed_shape).astype(target_dtype)
            input_dict[name] = data

        return input_dict
    except Exception as e:
        logger.error(e)
        return None


def run_inference(model_proto_bytes):
    """
    运行推理。改进版：优先从 Session 获取输入要求，实现类型自适应。
    """
    try:
        sess = ort.InferenceSession(model_proto_bytes, providers=['CPUExecutionProvider'])

        input_dict = {}
        # 1. 动态获取 Session 需要的输入
        for input_meta in sess.get_inputs():
            name = input_meta.name
            shape = input_meta.shape
            dtype_str = input_meta.type  # e.g. 'tensor(float16)'

            # 处理动态 Shape (None 或 str)
            fixed_shape = []
            for s in shape:
                if isinstance(s, int) and s > 0:
                    fixed_shape.append(s)
                else:
                    # 如果用户在 O0 没指定 shapes，这里兜底用 1
                    fixed_shape.append(1)

            # 类型映射
            if "float16" in dtype_str:
                target_dtype = np.float16
            elif "int64" in dtype_str:
                target_dtype = np.int64
            else:
                target_dtype = np.float32

            # 生成随机数据
            data = np.random.randn(*fixed_shape).astype(target_dtype)
            input_dict[name] = data

        outputs = sess.run(None, input_dict)
        return outputs, None
    except Exception as e:
        logger.error(e)
        return None, str(e)


def calculate_rel_error(original_proto: ModelStats, optimized_proto: ModelStats):
    """
    计算相对误差。不再依赖外部传入的 inputs_info，实现全自动收敛。
    """
    # 转换 bytes
    orig_bytes = original_proto.model.serialize_to_string()
    opt_bytes = optimized_proto.model.serialize_to_string()

    # 1. 运行原始模型
    ref_outs, err_ref = run_inference(orig_bytes)
    if err_ref:
        return "ERR_TYPE"

    # 2. 运行优化模型
    opt_outs, err_opt = run_inference(opt_bytes)
    if err_opt:
        return "ERR_TYPE"

    # 3. 计算误差
    errors = []
    for r, o in zip(ref_outs, opt_outs):
        r_f32 = r.astype(np.float32)
        o_f32 = o.astype(np.float32)

        # 针对 FP16 计算余弦相似度（Cosine Similarity）比相对误差更科学
        inner_prod = np.sum(r_f32 * o_f32)
        iter_norm = np.linalg.norm(r_f32) * np.linalg.norm(o_f32)
        cosine_sim = inner_prod / (iter_norm + 1e-7)

        # 转换成 1 - sim 作为误差项
        errors.append(1.0 - cosine_sim)

    avg_error = np.mean(errors)
    return avg_error * 100


def calculate_speed(original_proto: ModelStats, optimized_proto: ModelStats):
    # 转换 bytes
    orig_bytes = original_proto.model.serialize_to_string()
    opt_bytes = optimized_proto.model.serialize_to_string()

    # 加载
    orig_model = ort.InferenceSession(orig_bytes, providers=['CPUExecutionProvider'])
    opt_model = ort.InferenceSession(opt_bytes, providers=['CPUExecutionProvider'])
    random_inputs = generate_random_inputs(orig_model)

    # 预热
    _ = orig_model.run(None, random_inputs)
    _ = opt_model.run(None, random_inputs)

    # 统计
    iters = 30
    t1 = time.time_ns()
    for _ in range(iters):
        _ = orig_model.run(None, random_inputs)
    t2 = time.time_ns()
    orig_time_it = (t2 - t1) / 1_000_000 / iters
    t1 = time.time_ns()
    for _ in range(iters):
        _ = opt_model.run(None, random_inputs)
    t2 = time.time_ns()
    opt_time_it = (t2 - t1) / 1_000_000 / iters

    return orig_time_it, opt_time_it


def print_comparison_table(original: ModelStats, optimized: ModelStats, elapsed_time: float):
    """
    基础优化对比表：展示算子变化明细
    """
    console = Console()
    table = Table(title="🐕 ConvDog 算子优化对比", header_style="bold magenta", border_style="cyan", box=box.ROUNDED)

    table.add_column("算子/属性 (Op/Attr)", justify="right", style="bold")
    table.add_column("原始 (Original)", justify="center")
    table.add_column("优化后 (Optimized)", justify="center")

    # 1. 基础信息
    table.add_row("Model Name", os.path.basename(original.name), os.path.basename(optimized.name))
    table.add_row("Op Set / IR", f"v{original.opset} / {original.ir_version}",
                  f"v{optimized.opset} / {optimized.ir_version}")

    # 2. IO 信息 (取第一个作为代表)
    table.add_section()
    table.add_row("Input Info", original.inputs[0] if original.inputs else "N/A",
                  optimized.inputs[0] if optimized.inputs else "N/A")
    table.add_row("Output Info", original.outputs[0] if original.outputs else "N/A",
                  optimized.outputs[0] if optimized.outputs else "N/A")

    # 3. 算子列表 (带箭头对比)
    table.add_section()
    all_ops = sorted(list(set(original.op_counts.keys()) | set(optimized.op_counts.keys())))
    for op in all_ops:
        orig_cnt = original.op_counts.get(op, 0)
        opt_cnt = optimized.op_counts.get(op, 0)

        # 使用 diff 函数美化输出
        styled_opt = get_diff_str(orig_cnt, opt_cnt)
        table.add_row(op, str(orig_cnt), styled_opt)

    # 4. 模型体积
    table.add_section()
    table.add_row("Model Size", f"{original.size_mb:.2f} MB",
                  get_diff_str(original.size_mb, optimized.size_mb, is_size=True))

    # 5. 时间统计
    table.add_section()
    table.add_row("Elapsed Time", "-", f"[bold yellow]{elapsed_time:.3f} s[/]")

    console.print("\n")
    console.print(table)


def print_every_layer_quant_table(graph: ConvDogModel):
    """单独的量化状态表：展示权重类型分布"""
    console = Console()
    table = Table(title="💎 量化细节", header_style="bold blue")
    table.add_column("Weight Name", justify="left")
    table.add_column("Shape", justify="center")
    table.add_column("Precision", justify="center")

    for name, arr in graph.initializers.items():
        table.add_row(name, str(arr.shape), str(arr.dtype))

    console.print(table)


def print_quant_summary(original: ModelStats, optimized: ModelStats):
    """
    专家级看板：量化对比 + 回退分析 + 部署诊断
    """
    console = Console()

    # --- 1. 左侧：转换指标 ---
    comp_table = Table(box=box.SIMPLE, header_style="bold magenta")
    comp_table.add_column("Item", justify="right")
    comp_table.add_column("Change", justify="left")

    all_dts = sorted(list(set(original.weight_counts.keys()) | set(optimized.weight_counts.keys())))
    for dt in all_dts:
        o_c = original.weight_counts.get(dt, 0)
        n_c = optimized.weight_counts.get(dt, 0)
        if o_c != n_c:
            arrow = "[bold green]↑[/]" if n_c > o_c else "[bold red]↓[/]"
            comp_table.add_row(f"{dt} Weights", f"{o_c} -> {n_c} {arrow}")

    comp_table.add_row("Model Size",
                       f"{original.size_mb:.2f} -> {get_diff_str(original.size_mb, optimized.size_mb, is_size=True)}")

    # --- 2. 中间：回退层 (Fallback) ---
    fallback_layers = []
    for name, tensor in optimized.graph.tensors().items():
        if isinstance(tensor, gs.Constant):
            arr = tensor.values  # 获取底层的 numpy 数组
            if str(arr.dtype) == "float32":
                # 此时 arr 是 numpy 对象，拥有 .size 和 .itemsize
                mem_size_kb = (arr.size * arr.itemsize) / 1024
                fallback_layers.append((name, mem_size_kb))

    # 排序逻辑保持不变
    fallback_layers.sort(key=lambda x: x[1], reverse=True)

    fb_table = Table(box=box.SIMPLE, header_style="bold red")
    fb_table.add_column("Remaining FP32", max_width=25)
    fb_table.add_column("KB", justify="right")
    if not fallback_layers:
        fb_table.add_row("[green]Success: 0 fallback[/]", "")
    else:
        for name, size in fallback_layers[:8]:
            fb_table.add_row(name, f"{size:.0f}")

    # --- 3. 右侧：部署诊断 (Diagnostics) ---
    diag_table = Table(box=box.SIMPLE, header_style="bold cyan")
    diag_table.add_column("Metric", justify="right")
    diag_table.add_column("Status/Value", justify="left")

    # 对齐检查 (TensorCore 友好度)
    alignment_rate = (
                optimized.aligned_count / optimized.total_weight_layers * 100) if optimized.total_weight_layers > 0 else 0
    adj_status = "[green]Excellent[/]" if alignment_rate > 90 else "[yellow](Fair)[/]"
    diag_table.add_row("TC Alignment", f"{alignment_rate:.1f}% {adj_status}")

    # 算子消减率
    o_ops = sum(original.op_counts.values())
    n_ops = sum(optimized.op_counts.values())
    reduction = (1 - n_ops / o_ops) * 100 if o_ops > 0 else 0
    diag_table.add_row("Op Reduction", f"{reduction:.1f}%")

    # 误差分析
    rel_error = calculate_rel_error(original, optimized)
    if rel_error is not None:
        if rel_error == "ERR_TYPE":
            diag_table.add_row("Rel Error Δ", "[red bold]ORT 运行失败!!![/]")
        else:
            if rel_error < 0.8:
                color, status = "green", "(Perfect)"
            elif rel_error < 1.8:
                color, status = "yellow", "(Fair)"
            else:
                color, status = "red", "(Degraded)"
            diag_table.add_row("Rel Error Δ", f"[{color}]{rel_error:.3f}% {status}[/]")
    else:
        diag_table.add_row("Rel Error Δ", "[dim]ORT not found[/]")

    # 速度分析
    ori_speed, opt_speed = calculate_speed(original, optimized)
    speedup = ori_speed / (opt_speed + 1e-9)
    diag_table.add_row("Ori Speed", f"{ori_speed:.2f}ms")
    diag_table.add_row("Opt Speed", f"{opt_speed:.2f}ms")
    if speedup > 1.1:
        speed_color, speed_icon = "bold green", "🚀"
    elif speedup >= 0.95:
        speed_color, speed_icon = "bold yellow", "📈"
    else:
        speed_color, speed_icon = "bold red", "📉"
    diag_table.add_row("Speed Rate", f"{speed_icon} [{speed_color}]{speedup:.2f}[/]")

    # 精度分布条 (Visual Progress Bar)
    # 计算 FP16 占比
    fp16_ratio = optimized.dtype_params.get("float16", 0) / optimized.total_params if optimized.total_params > 0 else 0

    # 4. 组装
    main_group = Group(
        Columns([
            Panel(
                comp_table,
                title="📊 Precision Shift",
                border_style="magenta"
            ),
            Panel(fb_table, title="🚨 Fallback Details", border_style="red"),
            Panel(diag_table, title="🚀 Deployment Ready", border_style="cyan")
        ]),
        Panel(
            Text.assemble(
                ("Precision Distribution: ", "bold"),
                (f"FP16 {fp16_ratio * 100:.1f}% ", "yellow"),
                ("█" * int(fp16_ratio * 40), "yellow"),
                ("█" * (40 - int(fp16_ratio * 40)), "cyan"),
                (f" FP32 {(1 - fp16_ratio) * 100:.1f}%", "cyan"),
            ),
            border_style="dim",
            expand=False
        )
    )

    console.print(Panel(
        main_group,
        title=f"💎 ConvDog [bold]量化报告[/] - {os.path.basename(original.name)}",
        subtitle=f"Total: {original.total_params / 1e6:.2f}M Params | Efficiency: [bold green]{(1 - optimized.size_mb / original.size_mb) * 100:.1f}%[/]",
        border_style="bold blue", padding=(1, 2)
    ))
