import json
import sys

import numpy as np
import pandas as pd


def main(path: str) -> None:
    d = pd.read_csv(path)
    terms = [
        ("Attribute 1: Level 1", "Attribute 1", "Level 1"),
        ("Attribute 2: Level 1", "Attribute 2", "Level 1"),
        ("Attribute 3: Level 1", "Attribute 3", "Level 1"),
    ]
    x = np.column_stack([np.ones(len(d))] + [(d[attr] == level).astype(float).to_numpy() for _, attr, level in terms])
    y = d["chosen"].astype(float).to_numpy()
    bread_inv = np.linalg.inv(x.T @ x)
    beta = bread_inv @ x.T @ y
    resid = y - x @ beta
    meat = np.zeros((x.shape[1], x.shape[1]))

    for _, idx in d.groupby("response_id").groups.items():
        score = x[list(idx), :].T @ resid[list(idx)]
        meat += np.outer(score, score)

    vcov = bread_inv @ meat @ bread_inv
    clusters = d["response_id"].nunique()
    n = len(d)
    k = x.shape[1]
    vcov *= (clusters / (clusters - 1)) * ((n - 1) / (n - k))

    rows = []
    for i, term in enumerate(terms, start=1):
        rows.append({"term": term[0], "estimate": float(beta[i]), "se": float(np.sqrt(max(0, vcov[i, i])))})
    print(json.dumps(rows, indent=2))


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "sample-conjoint-profile-data.csv")