import torch
from thop import profile

from networks.subter_LeNet import SubTer_LeNet, SubTer_LeNet_Autoencoder
from networks.subter_LeNet_rf import SubTer_Efficient_AE, SubTer_EfficientEncoder

# Configuration
LATENT_DIMS = [32, 64, 128, 256, 512, 768, 1024]
BATCH_SIZE = 1
INPUT_SHAPE = (BATCH_SIZE, 1, 32, 2048)


def count_parameters(model, input_shape):
    """Count MACs and parameters for a model."""
    model.eval()
    with torch.no_grad():
        input_tensor = torch.randn(input_shape)
        macs, params = profile(model, inputs=(input_tensor,))
    return {"MACs": macs, "Parameters": params}


def format_number(num: float) -> str:
    """Format large numbers with K, M, B, T suffixes."""
    for unit in ["", "K", "M", "B", "T"]:
        if abs(num) < 1000.0 or unit == "T":
            return f"{num:3.2f}{unit}"
        num /= 1000.0


def main():
    # Collect results per latent dimension
    results = {}  # dim -> dict of 8 values
    for dim in LATENT_DIMS:
        # Instantiate models for this latent dim
        lenet_enc = SubTer_LeNet(rep_dim=dim)
        eff_enc = SubTer_EfficientEncoder(rep_dim=dim)
        lenet_ae = SubTer_LeNet_Autoencoder(rep_dim=dim)
        eff_ae = SubTer_Efficient_AE(rep_dim=dim)

        # Profile each
        lenet_enc_stats = count_parameters(lenet_enc, INPUT_SHAPE)
        eff_enc_stats = count_parameters(eff_enc, INPUT_SHAPE)
        lenet_ae_stats = count_parameters(lenet_ae, INPUT_SHAPE)
        eff_ae_stats = count_parameters(eff_ae, INPUT_SHAPE)

        results[dim] = {
            "lenet_enc_params": format_number(lenet_enc_stats["Parameters"]),
            "lenet_enc_macs": format_number(lenet_enc_stats["MACs"]),
            "eff_enc_params": format_number(eff_enc_stats["Parameters"]),
            "eff_enc_macs": format_number(eff_enc_stats["MACs"]),
            "lenet_ae_params": format_number(lenet_ae_stats["Parameters"]),
            "lenet_ae_macs": format_number(lenet_ae_stats["MACs"]),
            "eff_ae_params": format_number(eff_ae_stats["Parameters"]),
            "eff_ae_macs": format_number(eff_ae_stats["MACs"]),
        }

    # Build LaTeX table with tabularx
    header = (
        "\\begin{table}[!ht]\n"
        "\\centering\n"
        "\\renewcommand{\\arraystretch}{1.15}\n"
        "\\begin{tabularx}{\\linewidth}{lXXXXXXXX}\n"
        "\\hline\n"
        " & \\multicolumn{4}{c}{\\textbf{Encoders}} & "
        "\\multicolumn{4}{c}{\\textbf{Autoencoders}} \\\\\n"
        "\\cline{2-9}\n"
        "\\textbf{Latent $z$} & "
        "\\textbf{LeNet Params} & \\textbf{LeNet MACs} & "
        "\\textbf{Eff. Params} & \\textbf{Eff. MACs} & "
        "\\textbf{LeNet Params} & \\textbf{LeNet MACs} & "
        "\\textbf{Eff. Params} & \\textbf{Eff. MACs} \\\\\n"
        "\\hline\n"
    )

    rows = []
    for dim in LATENT_DIMS:
        r = results[dim]
        row = (
            f"{dim} & "
            f"{r['lenet_enc_params']} & {r['lenet_enc_macs']} & "
            f"{r['eff_enc_params']} & {r['eff_enc_macs']} & "
            f"{r['lenet_ae_params']} & {r['lenet_ae_macs']} & "
            f"{r['eff_ae_params']} & {r['eff_ae_macs']} \\\\"
        )
        rows.append(row)

    footer = (
        "\\hline\n"
        "\\end{tabularx}\n"
        "\\caption{Parameter and MAC counts for SubTer variants across latent dimensionalities.}\n"
        "\\label{tab:subter_counts}\n"
        "\\end{table}\n"
    )

    latex_table = header + "\n".join(rows) + "\n" + footer

    print(latex_table)


if __name__ == "__main__":
    main()