import os

import anndata
import pandas as pd

SPATIAL_HELP = 'Directory containing 10x visium spatial data (tissue_hires_image.png, scalefactors_json.json, and tissue_positions_list.csv) ' + 'or a directory containing `image.png`, `positions.image.csv` with headers barcode, x, and y, and optionally `diameter.image.txt` containing spot diameter',


def __add_visium(adata, spatial_directory):
    scale_factors_path = os.path.join(spatial_directory, 'scalefactors_json.json')
    tissue_hires_image_path = os.path.join(spatial_directory, 'tissue_hires_image.png')
    tissue_positions_list_path = os.path.join(spatial_directory, 'tissue_positions_list.csv')
    is_visium = True
    for path in [scale_factors_path, tissue_hires_image_path, tissue_positions_list_path]:
        if not os.path.exists(path):
            is_visium = False
            break
    if is_visium:
        import json
        with open(os.path.join(spatial_directory, 'scalefactors_json.json'), 'rt') as f:
            scalefactors = json.load(f)
            # {"spot_diameter_fullres": 89.49502418224989, "tissue_hires_scalef": 0.17011142,
            # "fiducial_diameter_fullres": 144.56888521748058, "tissue_lowres_scalef": 0.051033426}
        # barcode, in_tissue, array_row, array_col, pxl_col_in_fullres, pxl_row_in_fullres
        positions = pd.read_csv(tissue_positions_list_path, header=None)
        positions.columns = [
                'barcode',
                'in_tissue',
                'array_row',
                'array_col',
                'pxl_col_in_fullres',
                'pxl_row_in_fullres',
        ]
        positions.index = positions['barcode']
        positions = positions.reindex(adata.obs.index)
        spatial_coords = positions[['pxl_row_in_fullres', 'pxl_col_in_fullres']].values
        adata.obsm['tissue_hires'] = spatial_coords * scalefactors['tissue_hires_scalef']
        adata.uns['images'] = [dict(name='tissue_hires', image=tissue_hires_image_path,
            spot_diameter=scalefactors['spot_diameter_fullres'] * scalefactors['tissue_hires_scalef'])]
        return True
    else:
        return False


def __add_generic_spatial(adata, spatial_directory):
    # positions.image.csv with barcode, x, y
    # image.png
    # diameter.image.txt (optional)
    image_extensions = set(['.png', '.jpeg', '.jpg'])
    found = False
    for f in os.listdir(spatial_directory):
        name, ext = os.path.splitext(f)
        ext = ext.lower()
        if ext in image_extensions:
            positions_path = os.path.join(spatial_directory, 'positions.' + name + '.csv')

            if os.path.exists(positions_path):
                diameter_path = os.path.join(spatial_directory, 'diameter.' + name + '.txt')
                spot_diameter = None
                if os.path.exists(diameter_path):
                    with open(diameter_path, 'rt') as diameter_in:
                        spot_diameter = float(diameter_in.readline().strip())

                positions = pd.read_csv(positions_path, index_col='barcode')
                if not pd.api.types.is_string_dtype(positions.index):
                    positions.index = positions.index.astype(str)
                positions = positions.reindex(adata.obs.index)
                adata.obsm[name] = positions[['x', 'y']].values
                adata.uns['images'] = [dict(name=name, image=os.path.join(spatial_directory, f),
                    spot_diameter=spot_diameter)]
                found = True
    return found


def add_spatial(adata, spatial_directory):
    if not __add_visium(adata, spatial_directory):
        return __add_generic_spatial(adata, spatial_directory)
    else:
        return True


def get_cell_type_genes(cell_type):
    all_genes = []

    for cell_type_markers in cell_type['markers']:
        all_genes += cell_type_markers['genes']

    for i in range(len(all_genes)):
        gene = all_genes[i]
        last_char = gene[len(gene) - 1]
        if last_char == '+' or last_char == '-':
            all_genes[i] = gene[:len(gene) - 1]
    return all_genes


# remove genes in dict that are not in dataset
def filter_markers(adata, marker_dict):
    for category_name in marker_dict:
        markers = marker_dict[category_name]
        for key in markers:
            prior_size = len(markers[key])
            filtered_markers = set(markers[key]).intersection(adata.var.index)
            if len(filtered_markers) != prior_size:
                markers[key] = list(filtered_markers)
    return marker_dict


def get_markers(marker_paths):
    marker_dict = {}
    import json

    for marker_path in marker_paths:
        markers = {}
        with open(marker_path, 'rt') as f:
            marker_json = json.load(f)
            if 'title' in marker_json:
                marker_dict[marker_json['title']] = markers
                for cell_type in marker_json['cell_types']:
                    markers[cell_type['name']] = get_cell_type_genes(cell_type)
                    if 'subtypes' in cell_type:
                        for cell_sub_type in cell_type['subtypes']['cell_types']:
                            markers[cell_sub_type['name']] = get_cell_type_genes(cell_sub_type)
            else:
                key = os.path.splitext(os.path.basename(marker_path))[0]
                markers = marker_dict.get(key, {})
                marker_dict[key] = markers
                markers.update(marker_json)
    return marker_dict


def read_star_fusion_file(input_csv: str):
    """Load data from a STAR Fusion TSV file

    Parameters
    -----------

    input_csv : `str`
        The TSV file, gzipped or not, containing the STAR Fusion results.

    Returns
    -------
    """
    df = pd.read_csv(input_csv, sep='\t')
    df = df.sort_values('JunctionReadCount', ascending=False)
    df['value'] = True
    # df['value'] = df['value'].astype('category')
    df = df.drop_duplicates(['Cell', '#FusionName'], keep='first')
    # m = df.pivot_table('value', 'Cell', '#FusionName', fill_value=0)
    m = df.pivot('Cell', '#FusionName', 'value')
    m.fillna(False, inplace=True)
    return anndata.AnnData(obs=m)
