"""gen_shakespeare.py — Prepare the LEAF Shakespeare dataset for training.
"""

import argparse
import os
import shutil
import subprocess
import sys


LEAF_DIR = os.path.normpath(
    os.path.join(os.path.dirname(os.path.abspath(__file__)), 'leaf', 'data', 'shakespeare')
)


def get_config():
    parser = argparse.ArgumentParser(description='Prepare Shakespeare dataset via LEAF.')
    parser.add_argument('--seed1', type=int, default=777,
                        help='Seed label for the output directory.')
    parser.add_argument('--sf',   type=float, default=1.0,
                        help='Sampling fraction (fraction of characters to keep).')
    parser.add_argument('--k',    type=int,   default=0,
                        help='Minimum samples per user (users below this are dropped).')
    parser.add_argument('--tf',   type=float, default=0.8,
                        help='Train fraction for per-user train/test split.')
    return parser.parse_args()


def check_raw_data():
    raw_txt = os.path.join(LEAF_DIR, 'data', 'raw_data', 'pg100.txt')
    if not os.path.isfile(raw_txt):
        print('ERROR: Raw Shakespeare text not found at:')
        print('  ', raw_txt)
        print()
        print('Please download it first:')
        print('  cd', LEAF_DIR)
        print('  bash preprocess/get_data.sh')
        sys.exit(1)


def run_leaf_preprocess(sf, k, tf, smplseed, spltseed):
    """Run LEAF's preprocess.sh with the given options."""
    cmd = [
        'bash', 'preprocess.sh',
        '-s', 'niid',
        '--sf', str(sf),
        '-k',  str(k),
        '-t',  'sample',
        '--tf', str(tf),
        '--smplseed', str(smplseed),
        '--spltseed',  str(spltseed),
    ]
    print('>>> Running LEAF preprocessing:')
    print('    ' + ' '.join(cmd))
    result = subprocess.run(cmd, cwd=LEAF_DIR, check=True)
    return result


def copy_to_dataset_dir(seed1):
    """Copy LEAF output into dataset/shakespeare/data_{seed1}/."""
    src_train = os.path.join(LEAF_DIR, 'data', 'train')
    src_test  = os.path.join(LEAF_DIR, 'data', 'test')

    cur_path  = os.path.abspath(os.path.dirname(__file__))
    dst_base  = os.path.join(cur_path, 'dataset', 'shakespeare',
                              'data_{}'.format(seed1))
    dst_train = os.path.join(dst_base, 'train')
    dst_test  = os.path.join(dst_base, 'test')

    for src, dst in [(src_train, dst_train), (src_test, dst_test)]:
        if not os.path.isdir(src):
            print('ERROR: LEAF output directory not found:', src)
            print('The preprocessing step may have failed.')
            sys.exit(1)
        if os.path.exists(dst):
            shutil.rmtree(dst)
        shutil.copytree(src, dst)
        print('>>> Copied {} -> {}'.format(src, dst))


if __name__ == '__main__':
    args = get_config()

    print('=== Shakespeare dataset preparation ===')
    print('seed1={}, sf={}, k={}, tf={}'.format(args.seed1, args.sf, args.k, args.tf))

    check_raw_data()
    run_leaf_preprocess(sf=args.sf, k=args.k, tf=args.tf,
                        smplseed=args.seed1, spltseed=args.seed1)
    copy_to_dataset_dir(args.seed1)

    print()
    print('=== Done. Dataset written to:')
    print('    dataset/shakespeare/data_{}/'.format(args.seed1))
    print()
    print('You can now run training with:')
    print('    bash run_experiments.sh shakespeare')
