# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [default]
#     language: python
#     name: python3
# ---

# # Summary
#
# Generate training and validation datasets.
#
# ----

# # Imports

# +
import os
import random
import yaml
from pathlib import Path

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
# -

pd.set_option("max_columns", 100)

# # Parameters

NOTEBOOK_PATH = Path('training_validation_test_split')
NOTEBOOK_PATH

OUTPUT_PATH = Path(os.getenv('OUTPUT_DIR', NOTEBOOK_PATH.name)).resolve()
OUTPUT_PATH.mkdir(parents=True, exist_ok=True)
OUTPUT_PATH

# # `DATAPKG`

DATAPKG = {}

DATAPKG['uniparc-domain-wstructure'] = (
    Path(os.environ['DATAPKG_OUTPUT_DIR'])
    .joinpath("uniparc-domain-wstructure", "master")
)

# # Load data

adjacency_matrix_path = (
    DATAPKG['uniparc-domain-wstructure']
    .joinpath("remove_duplicate_matches", "adjacency_matrix.parquet")
)
adjacency_matrix_path

# # Construct training / validation / test datasets

# ## All Gene3D domains

# +
GENE3D_DOMAINS = sorted(p.name for p in adjacency_matrix_path.glob("database_id=*"))

random.seed(42)
random.shuffle(GENE3D_DOMAINS)
GENE3D_DOMAINS[:3]
# -

with NOTEBOOK_PATH.joinpath("all_gene3d_domains.yaml").open("wt") as fout:
    yaml.dump(GENE3D_DOMAINS, fout, default_flow_style=False)

# ## Split into training / validation / test

# +
breakpoint1 = len(GENE3D_DOMAINS) * 3 // 4
print(breakpoint1)
breakpoint2 = len(GENE3D_DOMAINS) * 7 // 8
print(breakpoint2)

TRAINING_DOMAINS = GENE3D_DOMAINS[:breakpoint1]
VALIDATION_DOMAINS = GENE3D_DOMAINS[breakpoint1:breakpoint2]
TEST_DOMAINS = GENE3D_DOMAINS[breakpoint2:]

print('----')
print(len(GENE3D_DOMAINS))
print(len(TRAINING_DOMAINS))
print(len(VALIDATION_DOMAINS))
print(len(TEST_DOMAINS))
# -

assert (len(TRAINING_DOMAINS) + len(VALIDATION_DOMAINS) + len(TEST_DOMAINS)) == len(GENE3D_DOMAINS)

# +
with NOTEBOOK_PATH.joinpath("training_domains.yaml").open("wt") as fout:
    yaml.dump(TRAINING_DOMAINS, fout, default_flow_style=False)
    
with NOTEBOOK_PATH.joinpath("validation_domains.yaml").open("wt") as fout:
    yaml.dump(VALIDATION_DOMAINS, fout, default_flow_style=False)
    
with NOTEBOOK_PATH.joinpath("test_domains.yaml").open("wt") as fout:
    yaml.dump(TEST_DOMAINS, fout, default_flow_style=False)

StarcoderdataJupyterScriptsDedupFiltered