# -*-coding:utf-8-*-
# ------------------
# Librarys
# ------------------


import sqlalchemy as sa
import numpy as np
import pandas as pd
# ------------------
from pycas.utils.decorator import helpstring
from pycas.utils.decorator import AppendHelp


# ------------------
# Variables
# ------------------


__all__ = ['GFF']
_helpdoc = dict()


# ------------------
# Errors
# ------------------


class SQLEngineError(ValueError):
    pass


class GFFError(ValueError):
    pass


# ------------------
# Classes
# ------------------


# {{{

_helpdoc['GFF'] = helpstring(
    describe='',
    parameterdicts={
        'gfffile': [
            'string, or bool.',
            'File path of gff or gtf file. ' +
            'If read gff or gtf from file, it should be asigned the file path.'+
            ' If the gff or gtf should be read from database, ' +
            'it should be asigned "False".'
        ],
        'tablename': [
            'string',
            'tablename to fetch or store gff or gtf file in the database.'
        ],
        'db': [
            'sqlalchemy.engine.base.Engine',
            'sqlalchemy Engine of database.'
        ],
        'if_exists': [
            'string, ["fail", "replace", "append"]',
            'parameter of pandas.DataFrame.to_sql.',
            '    "fail", if table exists, do nothing.',
            '    "replace", if table exists, drop it, recreate it, and insert data.',
            '    "append", if table exists, insert data. Create if does not exist.'
        ],
        'header': [
            'string',
            'None, if there is no header of gff or gtf file. ' +
            '0, if the first line fo gff or gtf file.'
        ]
    }
)

# }}}
@AppendHelp(_helpdoc['GFF'], join='')
class GFF:
    # {{{

    '''
    Manage gff or gtf
    '''
    def __init__(self,
                 gfffile,
                 tablename,
                 db,
                 if_exists,
                 header=None):
        # {{{

        if header is None:
            self.__colnames = ['seqname', 'source', 'feature',
                               'start', 'end', 'score',
                               'strand', 'frame', 'attribute']
        self.gfffile = gfffile
        self.tablename = tablename
        if isinstance(db, sa.engine.Engine):
            self.__db = db
        else:
            raise SQLEngineError('Wrong sqlalchemy engine type')
        self.__if_exists = if_exists
        if gfffile is False:
            self.gff = pd.read_sql_table(
                self.tablename,
                con=self.__db
            )
        else:
            self.__format = self.gfffile.split('.')[-1]
            self.gff = self.__parsegff(
                self.gfffile,
                self.__format,
                self.__colnames
            )
            self.__format = self.gfffile.split('.')[-1]

        # }}}

    def __parsegff(self, filename, fileformat, columns):
        # {{{

        '''
        Parse gff3 or gtf file into pandas DataFrame.

        Parameters
        ----------
        filename: the file path of the gff3 or gtf file.
        fileformat: the format of file, should be "gtf" or "gff".
        columns: the column names of the file.

        Returns
        -------
        DataFrame stored parsed information from gff or gtf

        Examples
        --------
        >>> import pandas as pd
        >>> mygff = parsegff(
                './hsa.gtf',
                'gtf',
                ['seqname', 'source', 'feature', 'start', 'end',
                 'score', 'strand', 'frame', 'attribute']
            )

        '''

        if fileformat == 'gtf':
            # make different separation according to fileformat.
            sep1, sep2 = ' "', '";'
        elif fileformat == 'gff3':
            sep1, sep2 = '=', ';'
        data = pd.read_table(
            filename,
            sep='\t',
            header=None,
            names=columns,
            index_col=False,
            comment='#'
        )                           # read gtf/gff data
        attr_split = data[columns[-1]].apply(
            lambda row: [
                x.split(sep1)
                for x in row.split(sep2) if len(x) > 2
            ]
        )
        # split the attributes in attribute columns
        attr_dict = attr_split.apply(
            lambda row: dict(
                zip(
                    [x[0].strip() for x in row],
                    [x[1].strip() for x in row]
                )
            )
        )
        # make splited attributes into dicts
        attr_columns = attr_dict.apply(
            lambda row: list(row.keys())
        ).tolist()                  # get attr columns names
        attr_names = list(
            set([
                attr_columns[i][j] for i in range(len(attr_columns))
                for j in range(len(attr_columns[i]))
            ])
        )
        # get attr columns names
        attr = pd.DataFrame(
            dict(
                zip(
                    attr_names,
                    [
                        pd.Series(
                            [x[attr_name] if attr_name in x else np.NaN
                             for x in attr_dict]
                        ) for attr_name in attr_names
                    ]
                )
            )
        )                           # make attr columns
        data = data.join(
            attr
        )
        # link attr columns with annotation.
        return data

        # }}}

    def to_sql(self):
        # {{{ store gtf to database

        self.gff.to_sql(
            self.tablename,
            con=self.__db,
            if_exists=self.__if_exists
        )

        # }}}

    # {{{ docstring

    _helpdoc['GFF.get_gene'] = helpstring(
        describe='',
        parameterdicts={
            'gene_symbol': [
                'string, or list like objects',
                'Gene symbol to return.'
            ]
        },
        returns='DataFrame contains the information.'
    )

    # }}}
    @AppendHelp(_helpdoc['GFF.get_gene'], join='')
    def get_gene(self,
                 gene_symbol):
        # {{{

        '''
        Get the gff information by gene_symbol.
        '''

        if (
                ('gene_id' not in self.gff.columns)
                or ('transcript_id' not in self.gff.columns)
        ):
            raise GFFError(
                'No gene_id or transcript_id information in gff file.'
            )
        else:
            pass
        if not hasattr(self, 'gff'):
            raise GFFError('No gff/gtf found.')
        if isinstance(gene_symbol, str):
            genelist = [gene_symbol]
        elif hasattr(gene_symbol, '__getitem__'):
            genelist = list(gene_symbol)
        else:
            raise ValueError('Wrong imput for gene_symbol')
        return self.gff.ix[
            self.gff['gene_id'].isin(gene_symbol)
        ]

        # }}}

    # {{{ docstring

    _helpdoc['GFF.get_gene_chromosome'] = helpstring(
        describe='',
        parameterdicts={
            'gene_symbol': [
                'string, or list like objects',
                'Gene symbol to return.'
            ]
        },
        returns='Series contains the information of gene chromosome.'
    )

    # }}}
    @AppendHelp(_helpdoc['GFF.get_gene_chromosome'], join='')
    def get_gene_chromosome(self, gene_symbol):
        # {{{

        '''
        Return gene chromosome information.
        '''

        gene = self.get_gene(gene_symbol)
        return gene.groupby(
            ['gene_id', 'transcript_id']
        ).apply(lambda x: x['seqname'].unique()[0])

        # }}}

    # {{{ docstring

    _helpdoc['GFF.get_gene_start'] = helpstring(
        describe='',
        parameterdicts={
            'gene_symbol': [
                'string, or list like objects',
                'Gene symbol to return.'
            ]
        },
        returns='Series contains the information of gene start.'
    )

    # }}}
    @AppendHelp(_helpdoc['GFF.get_gene_start'], join='')
    def get_gene_start(self, gene_symbol):
        # {{{

        '''
        Return gene start information.
        '''

        gene = self.get_gene(gene_symbol)
        return gene.groupby(
            ['gene_id', 'transcript_id']
        )['start'].min()

        # }}}

    # {{{ docstring

    _helpdoc['GFF.get_gene_end'] = helpstring(
        describe='',
        parameterdicts={
            'gene_symbol': [
                'string, or list like objects',
                'Gene symbol to return.'
            ]
        },
        returns='Series contains the information of gene end.'
    )

    # }}}
    @AppendHelp(_helpdoc['GFF.get_gene_end'], join='')
    def get_gene_end(self, gene_symbol):
        # {{{

        '''
        Return gene end information.
        '''

        gene = self.get_gene(gene_symbol)
        return gene.groupby(
            ['gene_id', 'transcript_id']
        )['end'].max()

        # }}}

    # {{{ docstring

    _helpdoc['GFF.get_gene_strand'] = helpstring(
        describe='',
        parameterdicts={
            'gene_symbol': [
                'string, or list like objects',
                'Gene symbol to return.'
            ]
        },
        returns='Series contains the information of gene end.'
    )

    # }}}
    @AppendHelp(_helpdoc['GFF.get_gene_strand'], join='')
    def get_gene_strand(self, gene_symbol):
        # {{{

        '''
        Return gene end information.
        '''

        gene = self.get_gene(gene_symbol)
        return gene.groupby(
            ['gene_id', 'transcript_id']
        ).apply(lambda x: x['strand'].unique()[0])

        # }}}

    # {{{ docstring

    _helpdoc['GFF.get_gene_feature'] = helpstring(
        describe='',
        parameterdicts={
            'gene_symbol': [
                'string, or list like objects',
                'Gene symbol to return.'
            ],
            'feature': [
                'string',
                'gene feature, should be "exon", "CDS", '+
                '"start_codon", "stop_codon" or "intron".'
            ]
        },
        returns='DataFrame of gene features.'
    )

    # }}}
    @AppendHelp(_helpdoc['GFF.get_gene_feature'], join='')
    def get_gene_feature(self, gene_symbol, feature):
        # {{{ return gene features

        '''
        Get the specific information of certain gene.
        '''
        gene = self.get_gene(gene_symbol)
        gene.reset_index(drop=True, inplace=True)
        if feature in ['exon', 'CDS', 'start_codon', 'stop_codon']:
            result = gene.ix[
                gene['feature'] == feature
            ][['gene_id', 'transcript_id', 'seqname',
               'start', 'end', 'strand']]
            result.reset_index(drop=True, inplace=True)
        elif feature == 'intron':
            def exon2intron(x):
                geneid = x.gene_id[:-1]
                geneid.reset_index(
                    drop=True,
                    inplace=True
                )
                transcriptid = x.transcript_id[:-1]
                transcriptid.reset_index(
                    drop=True,
                    inplace=True
                )
                istrand = x.strand[:-1]
                istrand.reset_index(
                    drop=True,
                    inplace=True
                )
                iseq = x.seqname[:-1]
                iseq.reset_index(
                    drop=True,
                    inplace=True
                )
                istart = x.end[:-1]
                istart.reset_index(
                    drop=True,
                    inplace=True
                )
                iend = x.start[1:]
                iend.reset_index(
                    drop=True,
                    inplace=True
                )
                return pd.concat(
                    [
                        geneid, transcriptid, iseq,
                        istart, iend, istrand
                    ],
                    axis=1,
                    ignore_index=True
                )
            result = gene.ix[gene['feature'] == 'exon']\
                         .sort_values(
                ['start', 'end']
            ).groupby(
                ['transcript_id']
            ).apply(
                exon2intron
            )
            result = result.reset_index(
                drop=True
            )
            result.columns = [
                'gene_id', 'transcript_id', 'seqname',
                'start', 'end', 'strand'
            ]
        else:
            raise ValueError('Wrong feature input')
        return result

        # }}}

    # }}}

####################
