import json
import os.path as osp
from glob import glob

import pandas as pd

# 1. N - Normal
# 2. V - PVC (Premature ventricular contraction)
# 3. \ - PAB (Paced beat)
# 4. R - RBB (Right bundle branch)
# 5. L - LBB (Left bundle branch)
# 6. A - APB (Atrial premature beat)
# 7. ! - AFW (Ventricular flutter wave)
# 8. E - VEB (Ventricular escape beat)

classes = ["N", "V", "\\", "R", "L", "A", "!", "E"]
lead = "MLII"
extension = "png"  # or `npy` for 1D
data_path = osp.abspath("../data/*/*/*/*/*.{}".format(extension))
val_size = 0.1  # [0, 1]

output_path = "/".join(data_path.split("/")[:-5])
random_state = 7

if __name__ == "__main__":
    dataset = []
    files = glob(data_path)

    for file in glob(data_path):
        *_, name, lead, label, filename = file.split("/")
        dataset.append(
            {
                "name": name,
                "lead": lead,
                "label": label,
                "filename": osp.splitext(filename)[0],
                "path": file,
            },
        )

    data = pd.DataFrame(dataset)
    data = data[data["lead"] == lead]
    data = data[data["label"].isin(classes)]
    data = data.sample(frac=1, random_state=random_state)

    val_ids = []
    for cl in classes:
        val_ids.extend(
            data[data["label"] == cl]
            .sample(frac=val_size, random_state=random_state)
            .index,
        )

    val = data.loc[val_ids, :]
    train = data[~data.index.isin(val.index)]

    train.to_json(osp.join(output_path, "train.json"), orient="records")
    val.to_json(osp.join(output_path, "val.json"), orient="records")

    d = {}
    for label in train.label.unique():
        d[label] = len(d)

    with open(osp.join(output_path, "class-mapper.json"), "w") as file:
        file.write(json.dumps(d, indent=1))