#!/usr/bin/env python3

import json
import numpy as np
import pandas as pd
from collections import defaultdict
from cell_reader import octagon, get_genes
import argparse

# a temporary horror show whilst we get up and running

def cells_dict(adata):
    index = adata.obs.index
    cells_dict = {}
    # genes_pair_list = defaultdict(list)
    # xy_dict = defaultdict(list)

    # df_genes = adata.groupby(['cell', 'gene1']).size().reset_index(name='count')

    # for index, row in df_genes.iterrows():
    #     genes_pair_list[row['cell']].append((row['gene1'], row['count']))

    # for index, row in df.iterrows():
    #     xy_dict[row['cell']].append([row['x'], row['y']])
    import re

    for index, cell in enumerate(index):
        polys = adata.obsm['simpler'][index][1:-1]
        polys = polys.split("\n")
        poly_array = []
        for poly in polys:
            poly = poly.strip()
            poly = poly[1:-1]
            poly = re.sub('\s+', ' ', poly)
            try:    
                coords = [int(x) for x in poly.split(" ")]
                #print(coords)
                poly_array.append(coords)   
                
                cells_dict[cell] = {
                    "mappings": {},
                    "genes": {},
                    "xy": list(adata.obsm['spatial'][index]),
                    "factors": {},
                    "poly": poly_array,
                }    
            except ValueError as e:
                print(e)
                pass

    return cells_dict


def molecules_dict(df):
    molecules_dict = defaultdict(list)

    for index, row in df.iterrows():
        molecules_dict[row['gene1']].append([row['x'], row['y']])

    return molecules_dict


def image_dict():
    cloud_target = open('cloud_target.txt').read().strip()
    url = 'https://s3.amazonaws.com/{}/wang.images/info.json'.format(
        cloud_target
    )
    image_dict = {
        'MERrfish': {
            'sample': 1,
            'tileSource': url
        }
    }

    return image_dict


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Create JSON with cell metadata from MERmaid CSV.')
    parser.add_argument(
        '--h5ad_file', required=True,
        help='H5AD file from MERfish data.')
    parser.add_argument(
        '--cells_file', type=argparse.FileType('x'),
        help='Write the cell data to this file.')
    parser.add_argument(
        '--molecules_file', type=argparse.FileType('x'),
        help='Create JSON with molecule locations.')
    parser.add_argument(
        '--genes_file', type=argparse.FileType('x'),
        help='Write a list of genes to this file.')
    parser.add_argument(
        '--images_file', type=argparse.FileType('x'),
        help='JSON file which will include image location.')
    args = parser.parse_args()

    import scanpy as sc
    adata = sc.read(args.h5ad_file)

    metadata = cells_dict(adata)

    if args.cells_file:
        json.dump(metadata, args.cells_file, indent=1)
    if args.molecules_file:
        json.dump(molecules_dict(adata), args.molecules_file, indent=1)
    if args.genes_file:
        json.dump(get_genes(metadata), args.genes_file, indent=1)
    if args.images_file:
        json.dump(image_dict(), args.images_file, indent=1)