#!/usr/bin/env python3
import re
import sys
from typing import Dict, List

def parse_ncu_log(log_path):
    metrics = []
    with open(log_path, 'r') as f:
        content = f.read()
    
    sections = content.split('='*80)
    for section in sections:
        if 'CUDA kernel' in section.lower() or 'kernel_' in section.lower():
            block_metrics = {}
            block_metrics['name'] = extract_kernel_name(section)
            for line in section.split('\n'):
                match = re.match(r'\s*(\S+)\s*[:]\s*([\d,]+\.?\d*)\s*(\w*)', line)
                if match:
                    key, value, unit = match.groups()
                    try:
                        block_metrics[key] = float(value.replace(',', ''))
                    except ValueError:
                        block_metrics[key] = value
            if block_metrics.get('name'):
                metrics.append(block_metrics)
    return metrics

def extract_kernel_name(section):
    match = re.search(r'Kernel:\s*(\S+)', section)
    if match:
        return match.group(1)
    for line in section.split('\n'):
        if 'kernel_' in line.lower() or '__' in line:
            parts = line.strip().split()
            for p in parts:
                if 'kernel' in p.lower():
                    return p.rstrip(':')
    return None

def analyze_bottlenecks(metrics: List[Dict]) -> List[Dict]:
    bottlenecks = []
    for m in metrics:
        issues = []
        sm_throughput = m.get('sm__throughput.avg_pct_of_peak_sustained_elapsed', 0)
        if sm_throughput < 50:
            issues.append(f"Low SM throughput: {sm_throughput:.1f}%")
        mem_throughput = m.get('gpu__dram_throughput.avg_pct_of_peak_sustained_elapsed', 0)
        if mem_throughput > 80:
            issues.append(f"Memory bandwidth bound: {mem_throughput:.1f}%")
        l2_hit = m.get('l2__hit_rate.pct', 0)
        if l2_hit < 50:
            issues.append(f"Low L2 hit rate: {l2_hit:.1f}%")
        occupancy = m.get('sm__warps_active.avg_pct_of_peak_sustained_active', 0)
        if occupancy < 50:
            issues.append(f"Low occupancy: {occupancy:.1f}%")
        if m.get('sass__inst_executed_op_shared_ld.sum', 0) > 100000:
            issues.append("High shared memory loads")
        if m.get('sass__inst_executed_op_shared_st.sum', 0) > 100000:
            issues.append("High shared memory stores")
        if issues:
            bottlenecks.append({'kernel': m['name'], 'issues': issues})
    return bottlenecks

def print_report(metrics: List[Dict], bottlenecks: List[Dict]):
    print("\n" + "="*70)
    print("NCU ANALYSIS REPORT")
    print("="*70)
    for m in metrics:
        print(f"\nKernel: {m.get('name', 'unknown')}")
        print(f"  SM Throughput: {m.get('sm__throughput.avg_pct_of_peak_sustained_elapsed', 'N/A'):.1f}%")
        print(f"  DRAM Throughput: {m.get('gpu__dram_throughput.avg_pct_of_peak_sustained_elapsed', 'N/A'):.1f}%")
        print(f"  L2 Hit Rate: {m.get('l2__hit_rate.pct', 'N/A'):.1f}%")
        print(f"  Occupancy: {m.get('sm__warps_active.avg_pct_of_peak_sustained_active', 'N/A'):.1f}%")
        print(f"  Grid: {m.get('launch__grid_size', 'N/A')}")
        print(f"  Block: {m.get('launch__block_size', 'N/A')}")
        print(f"  Registers: {m.get('launch__registers_per_thread', 'N/A')}")
    if bottlenecks:
        print("\nBOTTLENECKS:")
        for b in bottlenecks:
            print(f"\n  {b['kernel']}:")
            for issue in b['issues']:
                print(f"    - {issue}")

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Usage: python analyze_ncu.py <ncu_output.log>")
        sys.exit(1)
    metrics = parse_ncu_log(sys.argv[1])
    bottlenecks = analyze_bottlenecks(metrics)
    print_report(metrics, bottlenecks)