# /user/bin/env python3
# -*- coding: utf-8 -*-
"""KnowledgeDataBase class."""

import re
import os
from pprint import pformat
import pkg_resources

import logging
import pandas as pd
import openpyxl

from lc18report_dev.aescipher import AESCipher
from lc18report_dev.exceptions import WrongKdbFormatError

logger = logging.getLogger(__name__)


class KnowledgeDataBase:
    """Class for parsing knowledge file and related functions."""

    def __init__(self, kdb_file=''):
        """Knowledge database init function.
        """
        self._aa_records = {}
        self._cds_records = {}
        self._fusion_records = {}
        self._gene_records = {}
        self._med_records = {}
        self.file = kdb_file
        self._set_kdb_file()

    def _set_kdb_file(self):
        kdb_file = self.file
        logger.info(f"_set_kdb_file: kdb_file {kdb_file}")
        if not os.path.isfile(kdb_file):
            raise FileNotFoundError("知识库kdb_server未找到")
        self._parse_kdb_file(kdb_file)

    @staticmethod
    def _match_type(var_str):
        if re.match(r'^\s*p\.\S+\s*$', var_str):
            return 'aa'
        if re.match(r'^\s*c\.\S+\s*$', var_str):
            return 'cds'
        return 'fusion'

    def _parse_kdb_file(self, kdb_file):
        logger.info(f'Parsing kdb file {kdb_file}')

        # 读取excel的每个sheet xlsx
        knowledge_pd = pd.read_excel(kdb_file, engine='openpyxl', sheet_name='注释')
        gene_df      = pd.read_excel(kdb_file, engine='openpyxl', sheet_name='基因')
        med_df       = pd.read_excel(kdb_file, engine='openpyxl', sheet_name='药物-肺癌')

        for record in knowledge_pd.to_dict('records'):
            try:
                gene = record['基因'].strip().lower()
                level = record['等级'].strip().lower()
                var_str = record['变异'].strip().lower()
                sensitive = re.sub(
                    r'(\)|）)\W*(\w)',
                    r'\g<1>%%\g<2>',
                    record['提示敏感'].strip()
                )
                record['提示敏感'] = sensitive.split(r'%%')
                resistance = re.sub(
                    r'(\)|）)\W*(\w)',
                    r'\g<1>%%\g<2>',
                    record['提示耐药'].strip()
                )
                record['提示耐药'] = resistance.split(r'%%')
            except Exception:
                raise WrongKdbFormatError()

            record_type = self._match_type(var_str)
            if record_type == 'aa':
                self._aa_records[(gene, var_str)] = record
                logger.debug('Found variant record: ' + pformat(var_str))
            elif record_type == 'cds':
                self._cds_records[(gene, var_str)] = record
                logger.debug('Found variant record: ' + pformat(var_str))
            else:
                self._fusion_records[var_str] = record
                logger.debug('Found fusion record: ' + pformat(var_str))

        for gene_record in gene_df.to_dict('records'):
            gene = gene_record['基因'].strip().lower()
            self._gene_records[gene] = gene_record

        for med_record in med_df.fillna('').to_dict('records'):
            name = med_record['药物名称'].strip().lower()
            self._med_records[name] = med_record

    def get_record(self, gene, aa, cds):
        """Get variant record from knowledge database.

        Args:
            aa (str): amino acid change
            cds (str): base change

        Returns:
            dict: knowledge database record
        """
        logger.debug(f'Getting variant of {gene}-{aa}-{cds} from kdb.')
        clean_gene = gene.strip().lower()
        clean_aa = aa.strip().lower()
        clean_cds = cds.strip().lower()
        record = None

        if (clean_gene, clean_aa) in self._aa_records.keys():
            record = self._aa_records[(clean_gene, clean_aa)]
        elif (clean_gene, clean_cds) in self._cds_records.keys():
            record = self._cds_records[(clean_gene, clean_cds)]
        if record is None:
            logger.warning(f'No variant record of {gene}-{aa}-{cds} found.')
            return None
        logger.debug(f'Variant record found:\n{pformat(record)}')
        return record

    def get_fusion_record(self, gene1, gene2):
        """Get fusion record from knowledge database.

        Args:
            gene1 (str): first gene of fusion variant
            gene2 (str): second gene of fusion variant

        Returns:
            dict: knowledge database record
        """
        logger.info(f'Getting fusion of {gene1}, {gene2} from kdb.')
        clean_gene1 = gene1.strip().lower()
        clean_gene2 = gene2.strip().lower()

        try_keys = {f'{clean_gene1}-{clean_gene2}',
                    f'{clean_gene2}-{clean_gene1}'}
        matched = list(try_keys & self._fusion_records.keys())

        if not matched:
            logger.warning(f'No fusion record of {gene1}, {gene2} found.')
            return None

        fusion_record = self._fusion_records[matched[0]]
        logger.debug(f'Fusion record found:\n{pformat(fusion_record)}')
        return fusion_record

    def get_gene(self, gene: str):
        """Get gene record from knowledge database.

        Args:
            gene (str): gene name.

        Returns:
            dict: gene kdb record as dict type.
        """
        logger.debug(f'Getting gene of {gene} from kdb.')
        try:
            return self._gene_records[gene.strip().lower()]
        except KeyError:
            logger.warning(f'No gene record of {gene} found.')
            return None

    def get_med(self, med: str):
        """Get medicine record from knowledge database.

        Args:
            medicine (str): medicine name.

        Returns:
            dict: medicine kdb record as dict type.
        """
        logger.debug(f'Getting medicine of {med} from kdb.')
        try:
            return self._med_records[med.strip().lower()]
        except KeyError:
            logger.warning(f'No medicine record of {med} found.')
            return None
