import cooler
import pandas as pd
import numpy as np
import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed

def calculate_valid_bins_percentage(file_path, threshold=1000, num_threads=1, chunk_size=100000, resolutions=None):
    if file_path.endswith('.mcool'):
        return calculate_mcool_valid_bins_percentage(file_path, threshold, num_threads, chunk_size, resolutions)
    else:
        return calculate_cool_valid_bins_percentage(file_path, threshold, chunk_size)

def calculate_cool_valid_bins_percentage(file_path, threshold=1000, chunk_size=100000):
    clr = cooler.Cooler(file_path)
    return {
        'File': file_path,
        **calculate_bin_statistics(clr, threshold, chunk_size)
    }

def iter_bin_contact_sums(clr, chunk_size=100000):
    """Yield raw contact coverage, counting each matrix row exactly once.

    Cooler reconstructs symmetric-upper storage automatically. Read every
    column, including contacts outside the row batch, and count the diagonal
    once. Bound internal batches to keep sparse selections manageable.
    """
    if chunk_size <= 0:
        raise ValueError('chunk_size must be positive')
    row_batch_size = min(chunk_size, 4096)
    total_bins = clr.info['nbins']
    selector = clr.matrix(balance=False, sparse=True)
    for start in range(0, total_bins, row_batch_size):
        end = min(start + row_batch_size, total_bins)
        matrix = selector[start:end, :]
        # Accumulate fractional input counts in float64, not float32.
        yield np.asarray(matrix.sum(axis=1, dtype=np.float64)).ravel()

def calculate_bin_statistics(clr, threshold=1000, chunk_size=100000):
    total_bins = clr.info['nbins']
    total_valid_bins = sum(
        int(np.count_nonzero(bin_sums >= threshold))
        for bin_sums in iter_bin_contact_sums(clr, chunk_size)
    )
    return {
        'Valid Bin Percentage': (total_valid_bins / total_bins) * 100 if total_bins else 0,
        'Total Valid Bins': total_valid_bins,
        'Total Bins': total_bins
    }

def calculate_mcool_valid_bins_percentage(file_path, threshold=1000, num_threads=1, chunk_size=100000, resolutions=None):
    available_resolutions = cooler.fileops.list_coolers(file_path)

    print(f"Available resolutions in {file_path}: {available_resolutions}")
    if resolutions:
        resolutions = [f'/resolutions/{res}' for res in resolutions.split(',') if f'/resolutions/{res}' in available_resolutions]
        print(resolutions)
        if not resolutions:
            raise ValueError(f"No matching resolutions found in the .mcool file for specified resolutions.")
    else:
        resolutions = available_resolutions

    results = []

    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        future_to_resolution = {
            executor.submit(process_resolution, file_path, resolution_path, threshold, chunk_size): resolution_path
            for resolution_path in resolutions
        }

        for future in as_completed(future_to_resolution):
            resolution_path = future_to_resolution[future]
            try:
                result = future.result()
                if result is not None:
                    results.append(result)
            except Exception as e:
                print(f"Error processing resolution {resolution_path}: {e}")

    return {
        'File': file_path,
        'Results': results
    }

def process_resolution(file_path, resolution_path, threshold, chunk_size):
    clr_res = cooler.Cooler(f"{file_path}::{resolution_path}")

    print(f"Processing resolution: {resolution_path}")
    return {
        'Resolution': resolution_path.split('/')[-1],
        **calculate_bin_statistics(clr_res, threshold, chunk_size)
    }

def save_results_to_csv(results, output_path):
    if 'Results' in results:
        df = pd.DataFrame(results['Results'])
        df.insert(0, 'File', results['File'])
    else:
        df = pd.DataFrame([results])

    df.to_csv(output_path, index=False)
    print(f"Results saved to {output_path}")

def main():
    parser = argparse.ArgumentParser(description="Calculate valid bins percentage from .cool or .mcool files.")
    parser.add_argument("file_path", type=str, help="Path to the .cool or .mcool file")
    parser.add_argument("--output", type=str, required=True, help="Path to save the results as CSV file")
    parser.add_argument("--threads", type=int, default=1, help="Number of threads to use for processing")
    parser.add_argument("--chunk_size", type=int, default=100000, help="Maximum rows per sparse batch (internally capped at 4096; does not affect counts)")
    parser.add_argument("--resolutions", type=str, help="Comma-separated list of resolutions to process in the .mcool file")

    args = parser.parse_args()

    try:
        results = calculate_valid_bins_percentage(args.file_path, num_threads=args.threads, chunk_size=args.chunk_size, resolutions=args.resolutions)
        save_results_to_csv(results, args.output)
    except Exception as e:
        print(f"Error processing file {args.file_path}: {e}")

if __name__ == "__main__":
    main()