# File: entropy.py
# Copyright (C) 2026 Erick Ahmed
# SPDX-License-Identifier: AGPL-3.0-or-later

import argparse
from pathlib import Path

import numpy as np
import pandas as pd
import plotly.graph_objects as go

from utils.extractor import load_data
from utils.extractor import to_int

def _format_can_id(x):
    """Safely cleans CAN ID strings without altering their length or value."""
    if pd.isna(x):
        return "UNKNOWN"

    s = str(x).strip()
    if not s:
        return "UNKNOWN"

    if s.lower().startswith('0x'):
        s = s[2:]

    return s.upper()

def calculate_byte_entropy(df: pd.DataFrame) -> pd.DataFrame:
    """Calculates Shannon entropy per byte position for each identifier."""
    byte_cols = [f"b{i}" for i in range(8)]
    available_cols = [col for col in byte_cols if col in df.columns]
    if not available_cols:
        raise ValueError("No byte columns (b0-b7) found in the DataFrame")

    can_id_col = 'ID' if 'ID' in df.columns else 'Identifier'
    identifiers = df[can_id_col].apply(_format_can_id)

    df_bytes = df[available_cols].copy()
    for col in available_cols:
        df_bytes[col] = df_bytes[col].apply(to_int)

    def entropy(s: pd.Series) -> float:
        s = s.dropna()
        if s.empty:
            return 0.0
        p = s.value_counts(normalize=True)
        return -np.sum(p * np.log2(p))

    result = df_bytes.groupby(identifiers)[available_cols].agg(entropy)
    result.index.name = 'Identifier'
    return result


def plot_entropy_heatmap(entropy_df: pd.DataFrame, title: str) -> go.Figure:
    """Generates an interactive heatmap of byte-level Shannon entropy."""
    x = entropy_df.columns.tolist()
    y = entropy_df.index.tolist()
    z = entropy_df.values

    fig = go.Figure(
        data=go.Heatmap(
            z=z,
            x=x,
            y=y,
            colorscale=[
                [0.0, "#ffffff"],
                [0.15, "#fff7ec"],
                [0.35, "#fee8c8"],
                [0.55, "#fdd49e"],
                [0.75, "#fdbb84"],
                [1.0, "#ef6548"],
            ],
            xgap=3,
            ygap=3,
            text=np.round(z, 2),
            texttemplate="%{text}",
            textfont={
                "size": 11,
                "color": "#2a2a2a",
                "family": "Segoe UI, Arial, sans-serif",
            },
            hoverongaps=False,
            hovertemplate=(
                "<b>%{y}</b><br>"
                "Byte %{x}: %{z:.2f} bits<extra></extra>"
            ),
            colorbar=dict(
                title=dict(
                    text="Entropy (bits)",
                    side="top",
                    font=dict(size=13, color="#1a1a1a"),
                ),
                orientation="h",
                thickness=15,
                len=0.35,
                x=1.0,
                xanchor="right",
                y=1.02,
                yanchor="bottom",
                tickfont=dict(size=11, color="#2a2a2a"),
                tickformat=".1f",
                outlinewidth=0.5,
                outlinecolor="#cccccc",
            ),
        )
    )

    fig.update_layout(
        title=dict(
            text=title,
            font=dict(size=20, color="#1a1a1a"),
            x=0.5,
            xanchor="center",
            pad=dict(b=20),
        ),
        height=max(600, len(y) * 28 + 150),
        autosize=True,
        template="plotly_white",
        xaxis=dict(
            title=dict(text="Byte Position", font=dict(size=13, color="#1a1a1a")),
            side="top",
            dtick=1,
            showgrid=False,
            linecolor="#bdbdbd",
            tickfont=dict(size=12, color="#2a2a2a"),
            ticks="outside",
            ticklen=4,
            tickcolor="#cccccc",
        ),
        yaxis=dict(
            title=dict(text="PGN or CAN ID", font=dict(size=13, color="#1a1a1a")),
            autorange="reversed",
            showgrid=False,
            linecolor="#bdbdbd",
            tickfont=dict(size=12, color="#2a2a2a"),
            ticks="outside",
            ticklen=4,
            tickcolor="#cccccc",
            automargin=True,
        ),
        font=dict(family="Segoe UI, Arial, sans-serif", size=12, color="#2a2a2a"),
        hoverlabel=dict(
            bgcolor="white",
            font_size=13,
            font_family="Segoe UI",
            bordercolor="#cccccc",
        ),
        margin=dict(l=200, r=40, t=120, b=60),
    )
    return fig


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Analyze CAN bus byte-level entropy"
    )
    parser.add_argument(
        "input", type=Path, help="Path to the input CAN log file"
    )
    parser.add_argument(
        "output",
        type=Path,
        nargs="?",
        default=Path("entropy_report.html"),
        help="Path to the output HTML report",
    )
    parser.add_argument(
        "title",
        nargs="?",
        default="CAN Bus Byte-Level Entropy",
        help="Title for the HTML report",
    )
    args = parser.parse_args()

    df = load_data(args.input)
    entropy_df = calculate_byte_entropy(df)
    fig = plot_entropy_heatmap(entropy_df, title=args.title)

    config = {
        "responsive": True,
        "displaylogo": False,
        "scrollZoom": True,
        "modeBarButtonsToAdd": ["toggleSpikelines"],
        "toImageButtonOptions": {"format": "png", "scale": 2},
    }
    fig.write_html(str(args.output), include_plotlyjs="cdn", config=config)