import h5py
import hdf5plugin
import numpy as np

import argparse
import shutil
import os

parser = argparse.ArgumentParser(description="Drop regions on a given chromosome from shap/prediction h5s (in place, with a .bak backup)")
parser.add_argument("--h5_paths", type=str, required=True, help="comma-separated list of h5 paths to filter")
parser.add_argument("--exclude_chrom", type=str, default="chrM", help="chromosome to drop regions from")

args = parser.parse_args()

for h5_path in args.h5_paths.split(','):
    backup_path = h5_path + '.bak'

    with h5py.File(h5_path, 'r') as f:
        coords_chrom = f['coords_chrom'][()]
        coords_start = f['coords_start'][()]
        coords_end = f['coords_end'][()]
        hyp_scores = f['hyp_scores'][()]
        input_seqs = f['input_seqs'][()]

    keep_mask = coords_chrom != args.exclude_chrom.encode()
    num_dropped = (~keep_mask).sum()

    if num_dropped == 0:
        print(f'{h5_path}: no {args.exclude_chrom} regions found, skipping')
        continue

    coords_chrom = coords_chrom[keep_mask]
    coords_start = coords_start[keep_mask]
    coords_end = coords_end[keep_mask]
    hyp_scores = hyp_scores[keep_mask]
    input_seqs = input_seqs[keep_mask]
    num_examples = coords_chrom.shape[0]

    shutil.copy2(h5_path, backup_path)

    with h5py.File(h5_path, 'w') as f:
        coords_chrom_dset = f.create_dataset(
            "coords_chrom", (num_examples,),
            dtype=h5py.string_dtype(encoding="ascii"), **hdf5plugin.Blosc()
        )
        coords_chrom_dset[:] = coords_chrom

        coords_start_dset = f.create_dataset(
            "coords_start", (num_examples,), dtype=int, **hdf5plugin.Blosc()
        )
        coords_start_dset[:] = coords_start

        coords_end_dset = f.create_dataset(
            "coords_end", (num_examples,), dtype=int, **hdf5plugin.Blosc()
        )
        coords_end_dset[:] = coords_end

        hyp_scores_dset = f.create_dataset(
            "hyp_scores", hyp_scores.shape, **hdf5plugin.Blosc()
        )
        hyp_scores_dset[:, :, :] = hyp_scores

        input_seqs_dset = f.create_dataset(
            "input_seqs", input_seqs.shape, **hdf5plugin.Blosc()
        )
        input_seqs_dset[:, :, :] = input_seqs

    print(f'{h5_path}: dropped {num_dropped} {args.exclude_chrom} regions ({num_examples} remaining), backup saved to {backup_path}')