import pathlib

import joblib
import numcodecs
import numpy as np
import pandas as pd
import pyBigWig
import pyranges as pr
import zarr
from numba import jit


@jit(nopython=True)
def _get_chr_zarr_data(region_data, chr_size):
    zarr_data = np.zeros((chr_size,), dtype="int32")
    for start, end, value in region_data:
        zarr_data[start:end] = value
    return zarr_data


def coverage_to_zarr(coverage_path, zarr_path, chrom_sizes_path, chunk_size=20000000):
    """
    Convert a Tn5-currected cut site coverage file to a zarr file.

    Parameters
    ----------
    coverage_path : str
        Path to the Tn5-currected cut site coverage file generated by Genrich.
    zarr_path : str
        Path to the output zarr file.
    chrom_sizes_path : str
        Path to the chromosome sizes file.
    chunk_size : int
        Chunk size for the zarr file.

    Returns
    -------
    None
    """
    chrom_sizes = pd.read_csv(chrom_sizes_path, index_col=0, sep="\t", header=None).squeeze().to_dict()

    coverage = pd.read_csv(
        coverage_path,
        comment="#",
        sep="\t",
        usecols=[0, 1, 2, 3],
        dtype={"chr": "str", "start": "int32", "end": "int32", "experimental": "int32"},
    )
    coverage = coverage[coverage.iloc[:, 3] > 0].copy()

    pathlib.Path(zarr_path).mkdir(exist_ok=True, parents=True)
    root = zarr.group(store=zarr_path, overwrite=True)  # init zarr group
    zarr_compressor = numcodecs.Blosc(cname="zstd", clevel=3, shuffle=numcodecs.Blosc.SHUFFLE)  # Setup compressor

    for chr, chr_df in coverage.groupby("chr"):
        try:
            chr_size = chrom_sizes[chr]
        except KeyError:
            continue

        region_data = chr_df.iloc[:, 1:].values
        zarr_data = _get_chr_zarr_data(region_data=region_data, chr_size=chr_size)
        root.create_dataset(
            f"chrs/{chr}",
            data=zarr_data,
            chunks=chunk_size,
            compressor=zarr_compressor,
        )

    # final success flag
    pathlib.Path(f"{zarr_path}/.success").touch()
    return


def normalize_zarr(zarr_path, peak_path, chrom_sizes_path, output_zarr_path=None, chunk_size=20000000):
    """
    Normalize the raw cut site coverage zarr file.

    The normalization is done by subtracting the mean signal in non-peak regions and dividing by the 75th percentile of the signal in peak regions.
    Peak and non-peak regions are defined by a bed file, typically those ensambled peak list from ENCODE or atlas projects.

    Parameters
    ----------
    zarr_path : str
        Path to the input zarr file.
    peak_path : str
        Path to the peak regions bed file.
    chrom_sizes_path : str
        Path to the chromosome sizes file.
    output_zarr_path : str
        Path to the output zarr file. If None, only calculate the normalization factors and save them to the input zarr file directory at "{zar_path}/.norm".
    chunk_size : int
        Chunk size for the zarr file.

    Returns
    -------
    None
    """
    # Read peak and non-peak regions (bed file)
    genome_bed = pd.read_csv(chrom_sizes_path, header=None, sep="\t", names=["Chromosome", "End"])
    genome_bed["Start"] = 0
    genome_bed = pr.PyRanges(genome_bed[["Chromosome", "Start", "End"]])

    peaks = pr.read_bed(peak_path)
    non_peaks = genome_bed.subtract(peaks)
    peaks = peaks.as_df()
    non_peaks = non_peaks.as_df()

    # Read zarr track file to numpy array
    zarr_track_dir = zarr.open(zarr_path, mode="r")

    # Calculate the mean signal in non-peak regions
    non_peak_signals = []
    for chr_name, track in zarr_track_dir["chrs"].items():
        track = np.array(track)
        non_peaks_chr = non_peaks[non_peaks["Chromosome"] == chr_name]
        starts = non_peaks_chr.iloc[:, 1]
        ends = non_peaks_chr.iloc[:, 2]
        non_peak_signals_chr = [track[start:end].mean() for start, end in zip(starts, ends)]
        non_peak_signals += non_peak_signals_chr
    non_peak_mean = np.mean(non_peak_signals)

    # Calculate the 75th percentile signal in peak regions
    peak_signals = []
    for chr_name, track in zarr_track_dir["chrs"].items():
        track = np.array(track)
        peaks_chr = peaks[peaks["Chromosome"] == chr_name]
        starts = peaks_chr.iloc[:, 1]
        ends = peaks_chr.iloc[:, 2]
        peak_signals_chr = [track[start:end].mean() for start, end in zip(starts, ends)]
        peak_signals += peak_signals_chr
    peak_percentile = np.percentile(peak_signals, np.arange(99, 99.9, 0.1)).mean()

    joblib.dump(
        {"background_offset": non_peak_mean, "peak_scale": peak_percentile},
        f"{zarr_path}/.norm",
    )

    if output_zarr_path is not None:
        # Initialize the new zarr track similar to the zarr_track_dir
        new_zarr_track_dir = zarr.open_group(output_zarr_path, mode="w")
        new_zarr_track = new_zarr_track_dir.create_group("chrs")
        for chr_name, track in zarr_track_dir["chrs"].items():
            # Subtract the median signal in non-peak regions
            # Divide by the 75th percentile of the signal in peak regions
            new_data = (track[:] - non_peak_mean) / peak_percentile
            new_zarr_track.create_dataset(
                chr_name,
                data=new_data,
                chunks=chunk_size,
                shape=track.shape,
                dtype="float32",
            )


def zarr_to_bigwig(zarr_path, bigwig_path):
    """
    Convert a zarr file to a bigWig file.

    Parameters
    ----------
    zarr_path : str
        Path to the input zarr file.
    bigwig_path : str
        Path to the output bigWig file.

    Returns
    -------
    None
    """
    # Open the existing zarr group
    root = zarr.open_group(store=zarr_path, mode="r")
    chromsizes = {name: array.shape[0] for name, array in root["chrs"].items()}

    # Define headers
    header = list(chromsizes.items())

    # Open a new bigWig file for writing
    with pyBigWig.open(str(bigwig_path), "w") as bw:
        # Add the header
        bw.addHeader(header, maxZooms=10)

        # Loop over all chromosomes in the zarr group
        for chr_name, _ in header:
            print(f"Processing {chr_name}")
            # Load the dense array for this chromosome
            array = root["chrs"][chr_name][:]

            diff = np.diff(array, prepend=array[0] - 1, append=array[-1] + 1)
            change_points = np.where(diff)[0]

            starts = change_points[:-1]
            ends = change_points[1:]
            values = array[starts].astype(float)
            chroms = np.repeat(chr_name, len(starts))

            bw.addEntries(chroms, starts, ends, values=values)
