import os
from typing import Collection, List

import logging
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from lc18report_dev.sample import Sample
from lc18report_dev.base import base
from lc18report_dev.record_database import RecordDatabase
from lc18report_dev.template import Template
from lc18report_dev.exceptions import MissingResidentialIdError

logger = logging.getLogger(__name__)


class YueDong:

    MAX_TESTS = 6
    _db = RecordDatabase()

    def __init__(self):
        pass

    @classmethod
    def set_db_file(cls, db_file: str):
        cls._db.set_db_file(db_file)

    @staticmethod
    def save_sample(sample: Sample) -> None:
        if YueDong._db.get_sample_dict(sample.lib_id):
            YueDong._db.update_sample_rprt_dt(sample.get_sample_dict())
            logger.warning(
                f'Sample record {sample.lib_id} already exists in db.')
            return
        YueDong._db.save_sample(sample.get_sample_dict())
        YueDong._db.make_backup(os.path.join(os.getcwd(), 'yuedong_backup.db'))
        logger.info(f'Sample {sample.lib_id} saved')

    @staticmethod
    def _get_previous_samples(sample: Sample) -> Collection[Sample]:
        sample_dicts = YueDong._db.get_patient_sample_dicts(sample.resident_id)
        sample_dicts = list(sample_dicts)
        logger.debug(f'{len(sample_dicts)} samples found')
        samples = []
        for sample_dict in sample_dicts:
            if sample_dict['lib_id'] == sample.lib_id:
                continue
            samples.append(Sample.from_dict(sample_dict))
        if len(samples) > 5:
            samples = samples[:5]
        return samples

    @staticmethod
    def _get_yuedong_results(samples: List[Sample]) -> dict:
        brk = Template.BRK
        yuedong_results = {}
        gene_variants = {gene: {} for gene in Sample.genes}
        gene_fusions = {gene: {} for gene in Sample.genes}

        for i, sample in enumerate(samples):
            for variant in sample.variants:
                var_gene = variant['gene'].lower()
                aa = variant['aa'].strip()
                result = aa if aa else variant['cds'].strip()
                vaf = '> 80%'
                if variant['vaf'] <= 80:
                    vaf = '%.2f' % variant['vaf'] + '%'
                if result not in gene_variants[var_gene].keys():
                    gene_variants[var_gene][result] = {}
                gene_variants[var_gene][result][i] = vaf
            for fusion in sample.fusions:
                fus_gene = fusion['gene'].lower()
                result = fusion['cds'].strip()
                if result not in gene_fusions[fus_gene].keys():
                    gene_fusions[fus_gene][result] = {}
                gene_fusions[fus_gene][result][i] = 'N/A'

        for gene in Sample.genes:
            result_key = f'{gene}_results'
            for result, result_vafs in gene_variants[gene].items():
                if result_key not in yuedong_results.keys():
                    yuedong_results[result_key] = '#_red_#' + result
                else:
                    yuedong_results[result_key] += brk + result
            for result in gene_fusions[gene].keys():
                if result_key not in yuedong_results.keys():
                    yuedong_results[result_key] = '#_red_#' + result
                else:
                    yuedong_results[result_key] += brk + result
            if result_key not in yuedong_results.keys():
                yuedong_results[result_key] = '-'

        for i in range(YueDong.MAX_TESTS):
            for gene in Sample.genes:
                vaf_key = f'{gene}_vaf_{i + 1}'
                yuedong_results[vaf_key] = ''
                if i >= len(samples):
                    yuedong_results[vaf_key] = '-'
                    continue
                sample = samples[i]

                for result, result_vafs in gene_variants[gene].items():
                    if not yuedong_results[vaf_key]:
                        yuedong_results[vaf_key] = '#_red_#'
                    else:
                        yuedong_results[vaf_key] += brk
                    if i not in result_vafs.keys():
                        yuedong_results[vaf_key] += '未检出'
                    else:
                        yuedong_results[vaf_key] += result_vafs[i]

                for result, result_vafs in gene_fusions[gene].items():
                    if not yuedong_results[vaf_key]:
                        yuedong_results[vaf_key] = '#_red_#'
                    else:
                        yuedong_results[vaf_key] += brk
                    if i not in result_vafs.keys():
                        yuedong_results[vaf_key] += '未检出'
                    else:
                        yuedong_results[vaf_key] += result_vafs[i]

                if not yuedong_results[vaf_key]:
                    yuedong_results[vaf_key] = '-'

        return yuedong_results

    @staticmethod
    def get_yuedong_content(self, sample: Sample, img_file: str, force_mode=False) -> dict:
        if not sample.has_residential_id():
            sample.set_error(4)
            raise MissingResidentialIdError()
        samples = list(YueDong._get_previous_samples(sample))

        # 更新kdb_file为当前的kdb_file
        for prev_sample in samples:
            prev_sample.kdb_file = sample.kdb_file

        # 追加，最新一次的样本
        samples.append(sample)
        samples = sorted(samples)
        samples[-1]._reset_rprt_dt()
        vaf=[]
        # for sample in samples :
        # 仅计算当前样本的vaf最大值和平均值
        sample = samples[-1]
        if sample.variants:
            for var in sample.variants:
                vaf.append(var['vaf'])
        if vaf and len(vaf)!=0:
            if isinstance(vaf,list):
                self.max_vaf=max(vaf)
                self.avg_vaf=sum(vaf)/len(vaf)
            else:
                self.max_vaf=vaf
                self.avg_vaf=vaf
        else :
            self.max_vaf=0
            self.avg_vaf=0

        if self.max_vaf >= 80:
            self.max_vaf=80
        if self.avg_vaf >= 80:
            self.avg_vaf=80
        vaf.clear()

        for sample in samples :
            if sample!=samples[-1]:
                if sample.variants :
                    for var in sample.variants:
                        vaf.append(var['vaf'])
                if vaf and len(vaf):
                    max_sum=max(vaf)
                    avg_sum=sum(vaf)/len(vaf)
                else:
                    max_sum=0
                    avg_sum=0
                if max_sum >= 80:
                    max_sum=80
                if avg_sum>=80:
                    avg_sum=80
                sample.ct_ml='0' if pd.isnull(
                    sample.hGe*max_sum) or YueDong.ctdna_ml(
                        sample.hGe*max_sum)==0 else YueDong.ctdna_ml(sample.hGe*max_sum/100)
                sample.vm_ml='0' if pd.isnull(
                    sample.hGe*avg_sum) or YueDong.ctdna_ml(
                        sample.hGe*avg_sum)==0 else YueDong.ctdna_ml(sample.hGe*avg_sum/100)
                vaf.clear()
            else :
                if sample==samples[-1]:
                    sample.ct_ml='0' if pd.isnull(
                        samples[-1].hGe*self.max_vaf) or YueDong.ctdna_ml(
                            samples[-1].hGe*self.max_vaf)==0 else YueDong.ctdna_ml(samples[-1].hGe*self.max_vaf/100)
                    sample.vm_ml='0' if pd.isnull(
                        samples[-1].hGe*self.avg_vaf) or YueDong.ctdna_ml(
                            samples[-1].hGe*self.avg_vaf)==0 else YueDong.ctdna_ml(samples[-1].hGe*self.avg_vaf/100)
                    continue
                sample.ct_ml='0'
                sample.vm_ml='0'

        last_report_content = samples[-1].get_report_content(force_mode=force_mode)
        keys = last_report_content.keys()

        yuedong_content = {}
        yuedong_content['count'] = str(len(samples))
        for i in range(YueDong.MAX_TESTS):
            # Future tests
            if i >= len(samples):
                yuedong_content[f'count_{i + 1}'] = ''
                for key in keys:
                    new_key = key + f'_{i + 1}'
                    yuedong_content[new_key] = ''
                yuedong_content[f'rprt_dt_{i + 1}'] = '-'
                continue
            yuedong_content[f'count_{i + 1}'] = str(i + 1)
            sample_current: Sample = samples[i]
            logger.info(f"该患者还有 {sample_current.lib_id}")
            sample_content = sample_current.get_report_content(force_mode=force_mode)
            for key, value in sample_content.items():
                skip_suffixes = ['_exon', '_result', '_vaf']
                if any([key.endswith(suffix) for suffix in skip_suffixes]):
                    continue
                new_key = key + f'_{i + 1}'
                new_value = value if value else ' '
                if i + 1 == len(samples):
                    new_value = str(value) if str(value).startswith(
                        '#_red_#') else '#_red_#' + str(value)
                yuedong_content[new_key] = new_value

        yuedong_content[f'count_{len(samples)}'] = f'#_red_#{len(samples)}'
        yuedong_content.update(YueDong._get_yuedong_results(samples))

        # report content
        yuedong_content.update(last_report_content)

        # graph
        yuedong_content['hGe_graph'] = Template.IMG + img_file

        ba = base(
        base_info_file=self.base_info_file,record_xlsx=self.record_xlsx)
        yuedong_content=ba._parse_base_file(self.avg_vaf,self.max_vaf,samples,yuedong_content)
        try:
            if samples[-1].hGe > samples[-2].hGe:
                yuedong_content['graph_description'] = '本次检测发现ctDNA呈现上升趋势，应结合临床信息密切关注复发提示'
            else:
                yuedong_content['graph_description'] = '本次检测发现ctDNA呈现下降趋势，应结合临床其他信息评估复发风险'
        except IndexError:
            yuedong_content['graph_description'] = ''
        return yuedong_content

    @staticmethod
    def make_yuedong_diagram(self,samples: List[Sample], img_file: str) -> None:
        # [f'Test {i + 1}' for i in range(len(samples))]
        test_names =np.linspace(1,len(samples),len(samples))
        test_values = []
        for sample in samples :
            if sample.variants :
                vaf=[]
                for var in sample.variants:
                    vaf.append(var['vaf'])
                test_values.append(sample.hGe*max(vaf))
            else :
                test_values.append(sample.hGe*self.max_vaf)
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.set_ylim(0, max(test_values) + 200)
        ax.plot(
            test_names, test_values,
            linewidth=1,
            color='royalblue',
            solid_capstyle='round',
            marker='o',
        )
        for side in ['top', 'right', 'left']:
            ax.spines[side].set_visible(False)
        ax.grid(axis='y')
        ax.set_xlabel(samples[-1].id,fontsize=20)
        poly=np.polyfit(test_names,test_values,deg=1)
        z=np.polyval(poly,test_names)
        ax.plot(test_names,test_values,'o')
        ax.plot(test_names,z,label='y='+str(np.poly1d(poly)),linewidth=3)
        ax.legend(loc='upper left', bbox_to_anchor=(0, 1.0),fontsize=20)
        plt.savefig(img_file, bbox_inches='tight')

    def ctdna_ml(ml):
        if '.' not in str(ml) or ml==0 or ml==0.0:
            return str(int(ml))
        sum_right=len(str(ml).split(".")[1])
        sum_left=len(str(ml).split(".")[0])
        SIGN=0
        if '-' in str(ml):
            SIGN=1
        if sum_left>=4+SIGN :
            ml='%.0f' % float(ml)
        if sum_left==3+SIGN or (sum_left==1+SIGN and sum_right<2):
            ml='%.1f' % float(ml)
        if sum_left<=2+SIGN and sum_right>=2:
            ml='%.2f' % float(ml)
        return ml