import json
import pandas as pd
import pysam
import os 
import subprocess

from celescope.tools import utils
from celescope.tools.step import s_common, Step


class Summarize(Step):
    """
    ## Features

    - Convert 10X barcode of assemble result back to SGR barcode.

    - Generate Productive contigs sequences and annotation files.

    - Generate VDJ-annotation metrics in html.

    ## Output

    - `filtered_contig_annotations.csv` High-level annotations of each high-confidence contigs from cell-associated barcodes.

    - `filtered_contig.fasta` High-confidence contig sequences annotated in the filtered_contig_annotations.csv.

    - `productive_contig_annotations.csv` Annotations of each productive contigs from cell-associated barcodes. This is a subset of filtered_contig_annotations.csv.

    - `productive_contig.fasta` Productive contig sequences annotated in the productive_contig_annotations.csv.

    - `clonotypes.csv` High-level descriptions of each clonotype.

    """

    def __init__(self, args, display_title=None):
        Step.__init__(self, args, display_title=display_title)

        self.seqtype = args.seqtype
        self.version = os.path.dirname(args.soft_path).split('/')[-1].split('-')[-1]

        with open(args.barcode_convert_json, 'r') as f:
            self.tenX_sgr = json.load(f)

        # assemble result directory in 10X barcode.
        if args.not_refine:
            self.assemble_out = args.assemble_out
        else:
            self.assemble_out = args.assemble_out.replace('/outs', '_refine/outs')

        # in files
        annotation_file = f'{self.assemble_out}/filtered_contig_annotations.csv'
        self.df_annotation = pd.read_csv(annotation_file, sep=',', index_col=None)
        tenX_fasta_file = f'{self.assemble_out}/filtered_contig.fasta'
        self.tenX_fasta_fh = pysam.FastxFile(tenX_fasta_file)

        # out
        self.filtered_contig_annotations = f'{self.outdir}/filtered_contig_annotations.csv'
        self.filter_contig_fasta = f'{self.outdir}/filtered_contig.fasta'
    
    @utils.add_log
    def convert_barcode_to_SGR(self):
        """
        Convert 10X barcode to SGR barcode format.
        """

        self.df_annotation['barcode'] = self.df_annotation['barcode'].apply(lambda x: self.tenX_sgr[x.split('-')[0]])
        self.df_annotation['contig_id'] = self.df_annotation['contig_id'].apply(lambda x: 
            self.tenX_sgr[x.split('-')[0]]+'_'+x.split('_')[1]+'_'+x.split('_')[2])
        if int(self.version.split('.')[0]) < 4:
            self.df_annotation['productive'].replace({'True': True, 'None': False}, inplace=True)

        self.df_annotation.to_csv(self.filtered_contig_annotations, sep=',', index=False)
    

        with open(self.filter_contig_fasta, 'w') as f:
            for entry in self.tenX_fasta_fh:
                name = entry.name
                seq = entry.sequence
                attrs = name.split('_')
                new_name = self.tenX_sgr[attrs[0].split('-')[0]] + '_' + attrs[1] + '_' + attrs[2]
                f.write(f'>{new_name}\n{seq}\n')
        
        cmd = f'cp {self.assemble_out}/clonotypes.csv {self.outdir}'
        subprocess.check_call(cmd, shell=True)

    @staticmethod
    @utils.add_log
    def gen_productive_contig(annotation_file, fasta_file, outdir, prefix=''):
        """
        Generate productive_contig_annotation.csv and productive_contig.fasta file.
        """
        productive_contig = annotation_file[annotation_file['productive'] == True]
        productive_contig_id = set(productive_contig['contig_id'])
        productive_contig.to_csv(f'{outdir}/{prefix}productive_contig_annotations.csv', sep=',', index=False)

        productive_fasta = open(f'{outdir}/{prefix}productive_contig.fasta', 'w')
        with pysam.FastxFile(fasta_file, 'r') as f:
            for read in f:
                if read.name in productive_contig_id:
                    productive_fasta.write(">" + read.name + "\n" + read.sequence + "\n")
        productive_fasta.close()

    def run(self):
        self.convert_barcode_to_SGR()
        Summarize.gen_productive_contig(self.df_annotation, self.filter_contig_fasta, self.outdir)


def summarize(args):
    with Summarize(args) as runner:
        runner.run()


def get_opts_summarize(parser, sub_program):
    parser.add_argument('--seqtype', help='TCR or BCR', choices=['TCR', 'BCR'], required=True)
    parser.add_argument('--soft_path', help='soft path for cellranger')
    parser.add_argument("--not_refine", help="Do not perform the refine step. ", action="store_true")
    if sub_program:
        s_common(parser)
        parser.add_argument('--barcode_convert_json', help='json file', required=True)
        parser.add_argument('--assemble_out', help='directory of cellranger assemble result', required=True)
    return parser