#!/usr/bin/env python3

import argparse
import csv
import os
import random
import sys


def parse_args():
    parser = argparse.ArgumentParser(description='Stratified train/valid/test split with 7:1:2 ratio')

    parser.add_argument(
        '--input_csv',
        default='/home/qutaiping/nas/zhaocan/heart_diagnosis/diag_second_data.csv',
    )
    parser.add_argument(
        '--output_dir',
        default='/home/qutaiping/nas/zhaocan/heart_diagnosis/diag_data_split/',
    )
    parser.add_argument('--seed', type=int, default=42)

    return parser.parse_args()


def load_csv(csv_path):
    """Load CSV file and return header and rows."""
    with open(csv_path, 'r', newline='', encoding='utf-8-sig') as f:
        reader = csv.reader(f)

        header = next(reader)
        rows = list(reader)

    return header, rows


def calculate_split_sizes(total, ratios=(0.7, 0.1, 0.2)):
    """
    Calculate integer sample sizes according to the given ratios.

    Uses the largest remainder method to ensure that:
        train + valid + test = total
    """
    raw_sizes = [total * ratio for ratio in ratios]
    sizes = [int(size) for size in raw_sizes]
    remaining = total - sum(sizes)

    remainders = [(raw_sizes[i] - sizes[i], i) for i in range(len(ratios))]
    remainders.sort(reverse=True)

    for _, index in remainders[:remaining]:
        sizes[index] += 1

    return sizes


def allocate_by_class(class_counts, target_size, total_size):
    """
    Allocate target_size samples among classes according to their original class proportions.

    Uses the largest remainder method for integer allocation.
    """
    allocations = {}
    remainders = []

    for cls, count in class_counts.items():
        raw = count * target_size / total_size
        base = int(raw)
        allocations[cls] = base

        remainders.append((raw - base, cls))

    current_total = sum(allocations.values())
    remaining = target_size - current_total

    # Classes with larger fractional parts receive the remaining samples first.
    remainders.sort(key=lambda x: (-x[0], x[1]))

    for _, cls in remainders[:remaining]:
        allocations[cls] += 1

    return allocations


def stratified_split(rows, header, seed):
    """
    Perform stratified 7:1:2 train/valid/test split.
    """
    class_idx = header.index('class')

    # Group samples by class.
    by_class = {}

    for row in rows:
        cls = row[class_idx].strip()
        by_class.setdefault(cls, []).append(row)

    total_size = len(rows)

    if total_size == 0:
        raise ValueError('No samples available for splitting.')

    # Calculate total sizes of train/valid/test.
    train_size, valid_size, test_size = calculate_split_sizes(total_size)

    print('\n========== Dataset Size ==========')
    print(f'Total: {total_size}')
    print(
        f'Train: {train_size} '
        f'({train_size / total_size:.2%})'
    )
    print(
        f'Valid: {valid_size} '
        f'({valid_size / total_size:.2%})'
    )
    print(
        f'Test : {test_size} '
        f'({test_size / total_size:.2%})'
    )

    # Original number of samples in each class.
    class_counts = {cls: len(items) for cls, items in by_class.items()}

    train_counts = allocate_by_class(class_counts, train_size, total_size)

    remaining_counts = {cls: (class_counts[cls] - train_counts[cls]) for cls in class_counts}

    remaining_total = valid_size + test_size

    valid_counts = allocate_by_class(remaining_counts, valid_size,  remaining_total)

    test_counts = {cls: (remaining_counts[cls] - valid_counts[cls]) for cls in class_counts}

    print('\n========== Stratified Allocation ==========')
    print(
        f'{"Class":<15}'
        f'{"Total":>10}'
        f'{"Train":>10}'
        f'{"Valid":>10}'
        f'{"Test":>10}'
    )

    for cls in sorted(class_counts):
        print(
            f'{cls:<15}'
            f'{class_counts[cls]:>10}'
            f'{train_counts[cls]:>10}'
            f'{valid_counts[cls]:>10}'
            f'{test_counts[cls]:>10}'
        )

    # Randomly split samples within each class.
    rng = random.Random(seed)

    train_rows = []
    valid_rows = []
    test_rows = []

    for cls, items in by_class.items():
        items = items.copy()

        # Shuffle samples within each class.
        rng.shuffle(items)

        n_train = train_counts[cls]
        n_valid = valid_counts[cls]

        train_rows.extend(items[:n_train])
        valid_rows.extend(items[n_train:n_train + n_valid])
        test_rows.extend(items[n_train + n_valid:])

    # Shuffle each final dataset again so that different classes do not appear in blocks.
    rng.shuffle(train_rows)
    rng.shuffle(valid_rows)
    rng.shuffle(test_rows)

    return (train_rows, valid_rows, test_rows)


def write_csv(path, header, rows):

    with open(path, 'w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        writer.writerow(header)
        writer.writerows(rows)


def get_ids(rows, id_idx):

    return [row[id_idx].strip() for row in rows]


def calculate_distribution(rows, class_idx):
    """Calculate class counts and class distribution."""
    counts = {}

    for row in rows:
        cls = row[class_idx].strip()
        counts[cls] = counts.get(cls, 0) + 1

    total = len(rows)

    distribution = {cls: (count / total if total else 0) for cls, count in counts.items()}

    return counts, distribution


def main():
    args = parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    header, rows = load_csv(args.input_csv)

    class_idx = header.index('class')

    class_counts, class_dist = (calculate_distribution(rows, class_idx))

    print('\n========== Original Class Distribution ==========')

    print(
        f'{"Class":<15}'
        f'{"Count":>10}'
        f'{"Ratio":>12}'
    )

    for cls in sorted(class_counts):
        print(
            f'{cls:<15}'
            f'{class_counts[cls]:>10}'
            f'{class_dist[cls]:>11.2%}'
        )

    train_rows, valid_rows, test_rows = (stratified_split(rows, header, args.seed))

    train_path = os.path.join(args.output_dir, 'stage2_train.csv')
    valid_path = os.path.join(args.output_dir, 'stage2_valid.csv')
    test_path = os.path.join(args.output_dir, 'stage2_test.csv')

    write_csv(train_path, header, train_rows)
    write_csv(valid_path, header, valid_rows)
    write_csv(test_path, header, test_rows)


if __name__ == '__main__':
    main()