from typing import Dict
from pathlib import Path
import numpy as np
import pandas as pd
import pyranges as pr
from bamread import read_bam


cluster_col_order = [
    'Chromosome', 'Start', 'End', 'polyA_site', 'count', 'Strand',
    'Feature', 'gene_id', 'tpm', 'gene_count', 'usage',
    'fracA', 'signal', 'annotated_site'
]
"""List[str]: Columns of poly(A) clusters bed file."""

tss_cluster_col_order = [
    'Chromosome', 'Start', 'End', 'tss_site', 'count', 'Strand',
    'Feature', 'gene_id', 'tpm', 'gene_count', 'usage', 'annotated_site'
]
"""List[str]: Columns of tss clusters bed file."""


def read_polyA_cluster(path: str):
    '''
    Read poly(A) cluster file generated by LAPA.

    Args:
      path: Path to LAPA poly(A) cluster bed file.
    '''
    df = pd.read_csv(path, header=None, sep='\t')
    df.columns = cluster_col_order
    return df


def read_tss_cluster(path: str):
    '''
    Read tss cluster file generated by LAPA.

    Args:
      path: Path to LAPA TSS cluster bed file.
    '''
    df = pd.read_csv(path, header=None, sep='\t')
    df.columns = tss_cluster_col_order
    return df


def read_talon_read_annot(path: str, usecols=(
        'read_name', 'chrom', 'read_start', 'read_end',
        'strand', 'annot_gene_id', 'annot_transcript_id', 'dataset')):
    '''
    Read TALON read_annot file

    Args:
      path: Path to TALON read_annot.
      usecols: Only read subset of columns.
    '''
    df = pd.read_csv(path, sep='\t', usecols=usecols)
    df = df.rename(columns={
        'chrom': 'Chromosome',
        'read_start': 'Start',
        'read_end': 'End',
        'strand': 'Strand',
        'annot_gene_id': 'gene_id',
        'annot_transcript_id': 'transcript_id',
        'dataset': 'sample'
    })

    start = np.where(df['Start'] < df['End'], df['Start'], df['End'])
    end = np.where(df['Start'] > df['End'], df['Start'], df['End'])
    df['Start'] = start.copy()
    df['End'] = end.copy()

    return df


def read_talon_read_annot_three_prime_count(path: str):
    '''
    Read three_prime ends of TALON read_annot file.

    Args:
      path: Path to TALON read_annot.
    '''
    df = read_talon_read_annot(path)
    return _read_talon_read_annot_three_prime_count(df)


def _read_talon_read_annot_three_prime_count(df):

    df['End'] = np.where(df['Strand'] == '-', df['Start'], df['End'])
    return _read_talon_read_counts(df)


def read_talon_read_annot_five_prime_count(path: str):
    '''
    Read three_prime ends of TALON read_annot file.

    Args:
      path: Path to TALON read_annot.
    '''
    df = read_talon_read_annot(path)
    return _read_talon_read_annot_five_prime_count(df)


def _read_talon_read_annot_five_prime_count(df):
    df['End'] = np.where(df['Strand'] == '-', df['End'], df['Start'])
    return _read_talon_read_counts(df)


def _read_talon_read_counts(df):
    del df['Start']
    df['Start'] = df['End'] - 1
    df['count'] = 1

    return df[['Chromosome', 'Start', 'End', 'Strand', 'count', 'sample']]


def read_chrom_sizes(chrom_size_file: str) -> Dict:
    '''
    Read chrom_sizes file as dictionary of chroms as key and sizes as value.

    Args:
      chrom_size_file: Path to chromosome sizes file.
    '''
    df_chrom_sizes = pd.read_csv(chrom_size_file, sep='\t', header=None)
    chrom_sizes = df_chrom_sizes.set_index(0)[1].to_dict()
    return chrom_sizes


def bw_from_pyranges(gr: pr.PyRanges, value_col: str, chrom_sizes: str,
                     bw_pos_file: str, bw_neg_file: str):
    '''
    Save pyranges as bigwig file.

    Args:
      gr: pr.PyRanges object to save.
      value_col: Value column of the file to store.
      chrom_size_file: Path to chrom_sizes file
        or chrom_sizes as dictionary.
      bw_pos_file: Bigwig destination path to save positive strand scores.
      bw_neg_file: Bigwig destination path to save negative strand scores.
    '''
    if isinstance(chrom_sizes, str):
        chrom_sizes = read_chrom_sizes(chrom_sizes)

    gr = gr.subset(
        lambda df: df.Chromosome.isin(set(chrom_sizes.keys()))
        & (df.Start >= 0) # edge case occurs due to zero base indexing
    )
    chroms = set(gr.Chromosome.unique())
    chrom_sizes = {k: v for k, v in chrom_sizes.items() if k in chroms}

    if gr['+'].length > 0:
        gr['+'].to_bigwig(bw_pos_file, chromosome_sizes=chrom_sizes,
                          rpm=False, value_col=value_col)
    if gr['-'].length > 0:
        gr['-'].to_bigwig(bw_neg_file, chromosome_sizes=chrom_sizes,
                          rpm=False, value_col=value_col)
