# -*- coding: utf-8 -*-
__author__ = 'bars'

from collections import OrderedDict
from collections import defaultdict
from io import StringIO
import gzip
import pandas as pd
import sys


class Vcf():
    '''
    Main class for parsing vcf and its columns like sample format and variant info.
    '''

    def __init__(self, vcf_file, mode='file', compression='infer', chunksize=None):
        self.vcf_file = vcf_file
        self.mode = mode
        self.compression = compression
        # Chunksize not used right now.
        self.chunksize = chunksize
        # Generic VCF header without sample_names. For getting the sample_names.
        self.generic_header = ['CHROM',
                               'POS',
                               'ID',
                               'REF',
                               'ALT',
                               'QUAL',
                               'FILTER',
                               'INFO',
                               'FORMAT']

    def __enter__(self):
        if self.mode == 'file':
            self.get_header_n_meta_info()
            self.from_file()
            return self
        if self.mode == 'str':
            self.get_header_n_meta_info()
            self.from_string()
            return self

    def __exit__(self, exception_type, exception_value, traceback):
        pass

    def __str__(self):
        '''sample_names and number of variants'''
        return '{} {}'.format(' '.join(self.sample_names), len(self.vdf))

    def __repr__(self):
        '''sample_names and number of variants'''
        return '{} {} {}'.format(self.__class__.__name__, ' '.join(self.sample_names), len(self.vdf))

    def __len__(self):
        '''Number of variants'''
        return len(self.vdf)

    @property
    def sample_names(self):
        '''sample_names is header - generic_header'''
        return [colname for colname in self.header if colname not in self.generic_header]

    @property
    def vdf(self):
        '''These functions for setting vdf attribute.'''
        return self.variant_df

    @vdf.setter
    def vdf(self, new_df):
        '''These functions for setting vdf attribute.'''
        self.variant_df = new_df

    def get_header_n_meta_info(self):
        '''Get meta_info and header from vcf file.
           This function returns a dictionary.'''
        if self.mode == 'file':
            # Open file
            if self.vcf_file.endswith('.vcf.gz'):
                with gzip.open(self.vcf_file, 'rb') as f:
                    lines = f.readlines()
                # If gzip already pass as an argument pass
                # else set it to gzip
                if self.compression != 'infer':
                    pass
                else:
                    self.compression = 'gzip'
            elif self.vcf_file.endswith('.vcf'):
                with open(self.vcf_file, 'r') as f:
                    lines = f.readlines()
            else:
                sys.stderr.write('Are you sure this is a vcf file??')
                sys.exit()
        elif self.mode == 'str':
            lines = self.vcf_file.split('\n')
        # meta_info dict
        self.meta_info = OrderedDict()
        # Loop over lines
        for line in lines:
            if self.compression == 'gzip':
                line = line.decode('utf-8')
            # When it reaches the header break
            if '#CHROM' in line:
                if len(line.split('\t')) > 1:
                    sep = '\t'
                else:
                    sep = None
                self.header = [colname for colname in line.strip('#').replace(
                    '\n', '').split(sep)]
                break
            # Remove the pound and new line.  Split by equal sign.
            k, v = line.strip('#').replace('\n', '').split('=', 1)
            if k not in self.meta_info:
                self.meta_info[k] = [v]
            else:
                self.meta_info[k].append(v)

    def from_file(self):
        '''Reads the table. vdf stands for Variant Data Frame.
        comment='#' passes meta info and header.
        header is parsed from the file with header() function
        low_memory=False reads columns as object
        compression for gzip vcf files.
        chunksize pass'''
        self.variant_df = pd.read_csv(self.vcf_file,
                                      dtype=str,
                                      sep='\t',
                                      comment='#',
                                      delim_whitespace=True,
                                      names=self.header,
                                      low_memory=False,
                                      compression=self.compression)
        # chunks WIP
        # if self.chunksize is not None:
        #     for vdf_chunk in pd.read_csv(self.vcf_file,
        #                                    sep='\t',
        #                                    comment='#',
        #                                    names=self.header,
        #                                    low_memory=False,
        #                                    compression=compression,
        #                                    chunksize=self.chunksize):
        #         yield vdf_chunk

    def from_string(self):
        self.variant_df = pd.read_csv(StringIO(self.vcf_file),
                                      dtype=str,
                                      sep='\t',
                                      comment='#',
                                      delim_whitespace=True,
                                      names=self.header,
                                      low_memory=False)

    def get_variant_info(self, concat=False, drop=False):
        # Create key set from all rows
        all_key_set = set()
        for info_row in self.vdf['INFO']:
            # This line checks if the row is empty.
            # Fix for test vcf 'string_as_flag.vcf'
            if info_row == info_row:
                for variant_info in info_row.split(';'):
                    all_key_set.add(variant_info.split('=')[0])

        # Create info_dict that will be use to create info_df
        variant_info_dict = defaultdict(list)
        for index, info_row in zip(self.vdf.index, self.vdf['INFO']):
            variant_info_dict['index'].append(index)

            # If row is empty add NA to all keys and continue with the loop.
            if not info_row == info_row:
                for key in all_key_set:
                    variant_info_dict[key].append('NA')
                continue

            # If row is not empty collect its keys to another set.
            row_key_set = set()
            for variant_info in info_row.split(';'):
                row_key_set.add(variant_info.split('=')[0])

            # Add NA to missing keys.
            for key in all_key_set - row_key_set:
                variant_info_dict[key].append('NA')

            # Add values to dict
            for variant_info in info_row.split(';'):
                try:
                    row_key, v = variant_info.split('=')
                # If Its only a flag just add 1. E.g. ;NegativeTrainError;
                except ValueError:
                    row_key = variant_info.split('=')[0]
                    v = 1
                variant_info_dict[row_key].append(v)

        variant_info_df = pd.DataFrame.from_dict(variant_info_dict)
        variant_info_df.set_index('index', inplace=True)

        if concat:
            new_df = pd.concat([self.vdf, variant_info_df], axis=1)
            if drop:
                new_df.drop('INFO', axis=1, inplace=True)
            return new_df
        else:
            return variant_info_df

    def get_sample_format(self, sample_names=None, concat=False, drop=False):
        # If sample names is None self.sample_names is used.
        if sample_names is None:
            sample_names = self.sample_names
        else:
            sample_names = sample_names

        if len(sample_names) < 1:
            sys.exit('No sample column found.')

        samples_df_dict = {}
        for sample in sample_names:

            all_key_set = set()
            for format_row in self.vdf['FORMAT']:
                if format_row == format_row:
                    for sample_format in format_row.split(':'):
                        all_key_set.add(sample_format)

            sample_format_dict = defaultdict(list)
            for index, format_row, sample_row in zip(self.vdf.index, self.vdf['FORMAT'], self.vdf[sample]):
                sample_format_dict['index'].append(index)

                row_key_set = set()
                for sample_format in format_row.split(':'):
                    row_key_set.add(sample_format)

                for key in all_key_set - row_key_set:
                    sample_format_dict[key].append('NA')

                # Whole sample row missing. Fix for test vcf "strelka.vcf"
                if not sample_row == sample_row:
                    sample_row = 'NA'

                # Sample missing format. Fix for test vcf "example-4.1.vcf"
                while len(format_row.split(':')) > len(sample_row.split(':')):
                    sample_row += ':NA'

                for sample_format, sample_info in zip(format_row.split(':'), sample_row.split(':')):
                    sample_format_dict[sample_format].append(sample_info)

            sample_format_df = pd.DataFrame.from_dict(sample_format_dict)
            sample_format_df.set_index('index', inplace=True)
            samples_df_dict[sample] = sample_format_df

        if concat:
            for sample in samples_df_dict.keys():
                samples_df_dict[sample] = pd.concat(
                    [self.vdf, samples_df_dict[sample]], axis=1)

            if drop:
                if len(samples_df_dict) == 1:
                    samples_df_dict[sample_names[0]] = samples_df_dict[sample].drop(
                        sample_names[0], axis=1)
                else:
                    for sample in sample_names:
                        samples_df_dict[sample] = samples_df_dict[sample].drop(
                            sample_names, axis=1)

            if len(samples_df_dict) == 1:
                return samples_df_dict[list(samples_df_dict.keys())[0]]
            else:
                return samples_df_dict
        else:
            if len(samples_df_dict) == 1:
                return samples_df_dict[list(samples_df_dict.keys())[0]]
            else:
                return samples_df_dict

    def to_string(self):
        self.vcf_str = StringIO()
        for info_key, info_list in self.meta_info.items():
            for info in info_list:
                self.vcf_str.write('##{}={}\n'.format(info_key, info))
        self.vdf.rename(columns={'CHROM': '#CHROM'}, inplace=True)
        self.vdf.to_csv(self.vcf_str, sep='\t', index=False)
        return self.vcf_str

    def to_vcf(self, output_file):
        with open(output_file, 'w') as f:
            f.write(self.to_string().getvalue())

    def to_table(self, output_file):
        pass
