import csv
import os
import time

import torch
import torch.cuda.amp as amp
from loguru import logger


def _count_parameters(model):
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return total, trainable


def _cuda_available():
    return torch.cuda.is_available()


def _sync_cuda():
    if _cuda_available():
        torch.cuda.synchronize()


def _format_optional(value, precision=2):
    if value is None:
        return "NA"
    return f"{value:.{precision}f}"


def _append_csv(summary, csv_path):
    if not csv_path:
        return

    csv_dir = os.path.dirname(csv_path)
    if csv_dir:
        os.makedirs(csv_dir, exist_ok=True)

    fields = [
        "model_name",
        "parameters_m",
        "flops_g",
        "macs_g",
        "inference_time_ms",
        "trainable_parameters_m",
        "batch_size",
        "warmup",
        "repeat",
    ]
    write_header = not os.path.exists(csv_path) or os.path.getsize(csv_path) == 0
    with open(csv_path, "a", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fields)
        if write_header:
            writer.writeheader()
        writer.writerow({key: summary.get(key) for key in fields})


def _profile_flops(forward_fn, use_amp):
    try:
        from torch.profiler import ProfilerActivity, profile
    except Exception as exc:
        logger.warning(f"FLOPs profiling is unavailable: {exc}")
        return None

    activities = [ProfilerActivity.CPU]
    if _cuda_available():
        activities.append(ProfilerActivity.CUDA)

    try:
        with torch.no_grad(), profile(activities=activities, with_flops=True) as prof:
            with amp.autocast(enabled=use_amp and _cuda_available()):
                forward_fn()
        _sync_cuda()
    except Exception as exc:
        logger.warning(f"FLOPs profiling failed: {exc}")
        return None

    flops = sum((getattr(evt, "flops", 0) or 0) for evt in prof.key_averages())
    if flops <= 0:
        logger.warning(
            "FLOPs profiler returned 0. PyTorch only estimates FLOPs for supported ops."
        )
        return None
    return flops


def _profile_single_forward(forward_fn, use_amp, profile_flops):
    flops = None

    if not profile_flops:
        elapsed_ms = _measure_latency_ms(forward_fn, use_amp, warmup=0, repeat=1)
        return elapsed_ms, flops

    try:
        from torch.profiler import ProfilerActivity, profile
    except Exception as exc:
        logger.warning(f"FLOPs profiling is unavailable: {exc}")
        elapsed_ms = _measure_latency_ms(forward_fn, use_amp, warmup=0, repeat=1)
        return elapsed_ms, flops

    activities = [ProfilerActivity.CPU]
    if _cuda_available():
        activities.append(ProfilerActivity.CUDA)

    try:
        _sync_cuda()
        if _cuda_available():
            start = torch.cuda.Event(enable_timing=True)
            end = torch.cuda.Event(enable_timing=True)
        else:
            start_time = time.perf_counter()

        with torch.no_grad(), profile(activities=activities, with_flops=True) as prof:
            if _cuda_available():
                start.record()
            with amp.autocast(enabled=use_amp and _cuda_available()):
                forward_fn()
            if _cuda_available():
                end.record()

        _sync_cuda()
        if _cuda_available():
            elapsed_ms = start.elapsed_time(end)
        else:
            elapsed_ms = (time.perf_counter() - start_time) * 1000.0

        flops = sum((getattr(evt, "flops", 0) or 0) for evt in prof.key_averages())
        if flops <= 0:
            logger.warning(
                "FLOPs profiler returned 0. PyTorch only estimates FLOPs for supported ops."
            )
            flops = None
        return elapsed_ms, flops
    except Exception as exc:
        logger.warning(f"Single-forward profiling failed: {exc}")
        elapsed_ms = _measure_latency_ms(forward_fn, use_amp, warmup=0, repeat=1)
        return elapsed_ms, flops


def _measure_latency_ms(forward_fn, use_amp, warmup, repeat):
    warmup = max(int(warmup), 0)
    repeat = max(int(repeat), 1)

    with torch.no_grad():
        for _ in range(warmup):
            with amp.autocast(enabled=use_amp and _cuda_available()):
                forward_fn()
    _sync_cuda()

    if _cuda_available():
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        with torch.no_grad():
            for _ in range(repeat):
                with amp.autocast(enabled=use_amp and _cuda_available()):
                    forward_fn()
        end.record()
        _sync_cuda()
        return start.elapsed_time(end)

    start_time = time.perf_counter()
    with torch.no_grad():
        for _ in range(repeat):
            with amp.autocast(enabled=False):
                forward_fn()
    return (time.perf_counter() - start_time) * 1000.0


def profile_model_complexity(
    model,
    forward_fn,
    batch_size,
    model_name,
    warmup=10,
    repeat=50,
    use_amp=True,
    profile_flops=True,
    csv_path=None,
):
    model.eval()
    batch_size = max(int(batch_size), 1)
    warmup = max(int(warmup), 0)
    repeat = max(int(repeat), 1)

    total_params, trainable_params = _count_parameters(model)
    if warmup == 0 and repeat == 1:
        elapsed_ms, flops = _profile_single_forward(forward_fn, use_amp, profile_flops)
    else:
        elapsed_ms = _measure_latency_ms(forward_fn, use_amp, warmup, repeat)
        flops = None
    inference_time_ms = elapsed_ms / (repeat * batch_size)

    if profile_flops and flops is None and not (warmup == 0 and repeat == 1):
        flops = _profile_flops(forward_fn, use_amp)

    flops_g = flops / batch_size / 1e9 if flops is not None else None
    macs_g = flops_g / 2.0 if flops_g is not None else None

    summary = {
        "model_name": model_name,
        "parameters_m": total_params / 1e6,
        "flops_g": flops_g,
        "macs_g": macs_g,
        "inference_time_ms": inference_time_ms,
        "trainable_parameters_m": trainable_params / 1e6,
        "batch_size": batch_size,
        "warmup": warmup,
        "repeat": repeat,
    }

    logger.info("#########################################")
    logger.info("Model complexity and inference performance")
    logger.info(f"Model Name: {summary['model_name']}")
    logger.info(f"Parameters (M): {summary['parameters_m']:.2f}")
    logger.info(f"FLOPs (G, estimated): {_format_optional(summary['flops_g'])}")
    logger.info(f"MACs (G, estimated): {_format_optional(summary['macs_g'])}")
    logger.info(f"Inference Time (ms/sample): {summary['inference_time_ms']:.2f}")
    logger.info(f"Trainable Parameters (M): {summary['trainable_parameters_m']:.2f}")
    logger.info("#########################################")

    print(
        "COMPLEXITY_RESULT "
        f"model_name={summary['model_name']} "
        f"parameters_m={summary['parameters_m']:.2f} "
        f"flops_g={_format_optional(summary['flops_g'])} "
        f"macs_g={_format_optional(summary['macs_g'])} "
        f"inference_time_ms={summary['inference_time_ms']:.2f} "
        f"trainable_parameters_m={summary['trainable_parameters_m']:.2f}"
    )

    _append_csv(summary, csv_path)
    if csv_path:
        logger.info(f"Appended complexity result to: {csv_path}")

    return summary