# measure_all.py (加写log 'w', 增大num_points测真实, 加Bi mode调试)

import torch
import time
import numpy as np
import numpy.dtypes
import numpy
import matplotlib.pyplot as plt
from pointcept.engines.defaults import default_config_parser, default_setup
from pointcept.datasets import build_dataset, point_collate_fn
from pointcept.models import build_model
from pointcept.models.quantization.quant_utils import convert_ptv3_to_bi_ptv3
from pointcept.utils.config import DictAction

torch.serialization.add_safe_globals([numpy.dtypes.Float64DType, numpy.dtype, numpy.core.multiarray.scalar])

def flops_counter_hook(module, input, output):
    flops = 0
    if isinstance(module, torch.nn.Linear):
        batch_size = input[0].size(0)
        in_features = module.in_features
        out_features = module.out_features
        flops = batch_size * in_features * out_features * 2
    elif isinstance(module, torch.nn.Conv1d):
        batch_size = input[0].size(0)
        out_channels = module.out_channels
        in_channels = module.in_channels
        kernel_size = module.kernel_size[0]
        output_size = output.size(2)
        flops = batch_size * out_channels * in_channels * kernel_size * output_size * 2
    return flops

def measure_for_model(config_file, weight_path, quantize=False):
    print("="*40)
    print(f"Measuring for: {config_file.split('/')[-1]}, Weight: {weight_path}, Quantize: {quantize}")
    print("="*40)

    cfg = default_config_parser(config_file, None)
    cfg = default_setup(cfg)
    
    model = build_model(cfg.model)
    
    if quantize:
        print("INFO: Converting to Bi-PTV3...")
        model = convert_ptv3_to_bi_ptv3(model)
    
    print(f"INFO: Loading weights from {weight_path}...")
    weight = torch.load(weight_path, map_location="cpu", weights_only=True)
    if "state_dict" in weight:
        weight = weight["state_dict"]
    
    from collections import OrderedDict
    new_state_dict = OrderedDict()
    for k, v in weight.items():
        name = k[7:] if k.startswith('module.') else k
        new_state_dict[name] = v
    
    model.load_state_dict(new_state_dict, strict=False)
    
    model = model.cuda()
    model.eval()

    print("INFO: Preparing one sample for testing...")
    dataset = build_dataset(cfg.data.val)
    input_dict = dataset[0]
    
    num_points = 81920  # 正常大小测真实
    if 'coord' in input_dict:
        input_dict['coord'] = input_dict['coord'][:num_points]
    if 'feat' in input_dict:
        input_dict['feat'] = input_dict['feat'][:num_points]
    if 'segment' in input_dict:
        input_dict['segment'] = input_dict['segment'][:num_points]
    if 'origin_segment' in input_dict:
        input_dict['origin_segment'] = input_dict['origin_segment'][:num_points]
    if 'grid_coord' in input_dict:
        input_dict['grid_coord'] = input_dict['grid_coord'][:num_points]
    if 'inverse' in input_dict:
        input_dict['inverse'] = input_dict['inverse'][:num_points]
    if 'offset' in input_dict:
        input_dict['offset'] = torch.tensor([num_points], dtype=torch.int)
    
    input_dict = point_collate_fn([input_dict])
    for key in input_dict:
        if isinstance(input_dict[key], torch.Tensor):
            input_dict[key] = input_dict[key].cuda()
    
    total_flops = 0
    hooks = []
    def add_hooks(m):
        if isinstance(m, (torch.nn.Linear, torch.nn.Conv1d)):
            def hook_fn(module, input, output):
                nonlocal total_flops
                flops_added = flops_counter_hook(module, input, output)
                total_flops += flops_added
                print(f"DEBUG: Layer {module.__class__.__name__} FLOPs added: {flops_added}")
            hooks.append(m.register_forward_hook(hook_fn))
    
    model.apply(add_hooks)
    
    print("DEBUG: Starting model forward for FLOPs calculation...")
    with torch.no_grad():
        _ = model(input_dict)
    
    print("DEBUG: Model forward complete, total FLOPs calculated.")
    
    for h in hooks:
        h.remove()
    
    flops_g = total_flops / 1e9

    num_warmup = 100
    num_runs = 500
    
    print(f"INFO: Warming up for {num_warmup} iterations...")
    with torch.no_grad():
        for _ in range(num_warmup):
            _ = model(input_dict)
    
    torch.cuda.synchronize()

    timings = []
    peak_memory = 0
    print(f"INFO: Running measurement for {num_runs} iterations...")
    with torch.no_grad():
        for _ in range(num_runs):
            torch.cuda.synchronize()
            start_time = time.perf_counter()
            
            _ = model(input_dict)
            
            current_memory = torch.cuda.max_memory_allocated() / 1024**2  # MB
            peak_memory = max(peak_memory, current_memory)
            torch.cuda.reset_peak_memory_stats()
            
            torch.cuda.synchronize()
            end_time = time.perf_counter()
            
            timings.append((end_time - start_time) * 1000)

    avg_latency = np.mean(timings)
    std_latency = np.std(timings)

    print(f"Total FLOPs: {flops_g:.3f} GFLOPs")
    print(f"Average Latency: {avg_latency:.3f} ms")
    print(f"Peak Memory: {peak_memory:.2f} MB")

    return flops_g, avg_latency, peak_memory

def generate_plot(fp32_flops, fp32_latency, fp32_memory, short_flops, short_latency, short_memory, long_flops, long_latency, long_memory):
    models = ['FP32 PTV3', 'Bi-PTV3 Short QAT', 'Bi-PTV3 Long QAT']
    flops = [fp32_flops, short_flops, long_flops]
    latency = [fp32_latency, short_latency, long_latency]
    memory = [fp32_memory, short_memory, long_memory]

    fig, axs = plt.subplots(1, 3, figsize=(18, 5))
    
    # (a) Latency
    axs[0].bar(models, latency, color=['blue', 'orange', 'orange'])
    axs[0].set_title('Time Cost (ms)')
    axs[0].set_ylabel('ms')

    # (b) Memory
    axs[1].bar(models, memory, color=['blue', 'orange', 'orange'])
    axs[1].set_title('Storage Usage (MB)')
    axs[1].set_ylabel('MB')

    # (c) FLOPs vs Latency
    axs[2].scatter(latency, flops, color=['blue', 'orange', 'orange'])
    for i, txt in enumerate(models):
        axs[2].annotate(txt, (latency[i], flops[i]))
    axs[2].set_title('FLOPs (G) vs Time Cost (ms)')
    axs[2].set_xlabel('Time Cost (ms)')
    axs[2].set_ylabel('FLOPs (G)')

    plt.savefig('performance_showdown.png')
    print("Generated plot: performance_showdown.png")

def measure_all():
    models = [
        ("FP32 Baseline", "configs/s3dis/semseg-pt-v3m1-0-base.py", "exp/fp32_baseline/model/model_last.pth", False),
        ("Bi-PTV3 Short QAT", "configs/s3dis/semseg-pt-v3m1-0-base_qat.py", "exp/bi_ptv3_qat_sprint_run/model/model_last.pth", True),
        ("Bi-PTV3 Long QAT", "configs/s3dis/semseg-pt-v3m1-0-base_qat_long.py", "exp/bi_ptv3_qat_long_run/model/model_best.pth", True)
    ]

    results = []
    log_str = "Results Table:\n| Model | FLOPs (G) | Latency (ms) | Memory (MB) |\n|-------|-----------|--------------|-------------|\n"
    for name, config, weight, quantize in models:
        flops, latency, memory = measure_for_model(config, weight, quantize)
        results.append((name, flops, latency, memory))
        log_str += f"| {name} | {flops:.3f} | {latency:.3f} | {memory:.2f} |\n"

    # 写log文件
    with open('measure_results.log', 'w') as f:
        f.write(log_str)
    print("Results saved to measure_results.log")

    # 生成图
    generate_plot(results[0][1], results[0][2], results[0][3], results[1][1], results[1][2], results[1][3], results[2][1], results[2][2], results[2][3])

if __name__ == "__main__":
    measure_all()