# Copyright (c) Meta Platforms, Inc. and affiliates.
import sys
from pathlib import Path

import altair as alt
import pandas as pd
import pydantic
from omegaconf import OmegaConf


class ScalingPlotsConfig(pydantic.BaseModel):
    df_dir: str
    output_chart_dir: str
    frame_files: list[str]

    class Config:
        extra = "forbid"


def determine_family(key: str):
    if key.startswith("Megabyte++"):
        return "Megabyte++"
    elif key.startswith("BLT"):
        return "BLT"
    elif key.startswith("LLaMA"):
        return "LLaMA"
    elif key.startswith("Space"):
        return "Space"


file_to_vars = {}


def create_chart(df: pd.DataFrame, output_file: str):
    df["metric"] = df["bpb/not_heldout.jsonl"]
    df["family"] = df["key"].map(determine_family)
    model_domain = [
        "BLT Space ps=6",
        "BLT Space w/o cross-attn",
        "SpaceByte",
        "LLaMA 3 BPE",
        "Megabyte++ ps=4",
        "Megabyte++ ps=6",
    ]
    color_range = ["#1f77b4", "#1f77b4", "#1f77b4", "#ff7f0e", "#2ca02c", "#2ca02c"]
    shape_range = [
        "circle",
        "square",
        "cross",
        "diamond",
        "triangle-up",
        "triangle-down",
    ]
    color_scale = alt.Scale(domain=model_domain, range=color_range)
    shape_scale = alt.Scale(
        domain=model_domain,
        range=shape_range,
    )
    base_chart = alt.Chart(df).encode(
        x=alt.X("flops", title="Training FLOPS")
        .scale(type="log", domain=[2e20, 1.25e22])
        .axis(values=[2e20, 4e20, 8e20, 1e21, 2e21, 4e21, 8e21, 1e22]),
        y=alt.Y("metric", title="Bits per Byte (BPB)").scale(zero=False),
    )
    lines = base_chart.encode(
        color=alt.Color("key", title="Model Color", scale=color_scale, legend=None),
        strokeDash=alt.StrokeDash("family", title="Model Family", legend=None),
    ).mark_line()
    points = base_chart.encode(
        color=alt.Color("key", title="Model", scale=color_scale),
        shape=alt.Shape("key", title="", scale=shape_scale),
    ).mark_point(size=70)
    chart = (
        (lines + points)
        .resolve_scale(
            color="independent",
            shape="independent",
            # strokeDash="independent",
        )
        .configure_legend(orient="right")
        .properties(height=300, width=400)
    )
    print("Saving", output_file)
    chart.save(output_file)


def main():
    config_path = sys.argv[1]
    file_config = OmegaConf.load(config_path)
    # Omit program name and config file name
    cli_conf = OmegaConf.from_cli(sys.argv[2:])
    conf_dict = OmegaConf.to_container(
        OmegaConf.merge(file_config, cli_conf), resolve=True, throw_on_missing=True
    )
    plot_config = ScalingPlotsConfig(**conf_dict)
    df_dir = Path(plot_config.df_dir)
    chart_dir = Path(plot_config.output_chart_dir)
    chart_dir.mkdir(exist_ok=True, parents=True)
    for ff in plot_config.frame_files:
        path = df_dir / ff
        df = pd.read_json(path)
        print(df)
        print(df.columns)
        create_chart(df, chart_dir / f"{path.name}.pdf")


if __name__ == "__main__":
    main()