"""Excel spreadsheet IO functions"""

import contextlib
import logging
from copy import copy
from pathlib import Path
from typing import List, Optional, Set, Tuple, Dict, Union

import openpyxl
import pandas as pd
from xlsxwriter.workbook import Workbook
from xlsxwriter.worksheet import Worksheet

from xlavir.images import SheetImage
from xlavir.io.excel_sheet_dataframe import ExcelSheetDataFrame, SheetName
from xlavir.qc import QualityRequirements
from xlavir.util import get_col_widths, get_row_heights
from xlavir.__about__ import __version__

logger = logging.getLogger(__name__)


def copy_spreadsheet(src_path: Path,
                     dest_path: Path,
                     source_sheet_index: int = 0) -> None:
    """Copy a spreadsheet from a source Excel spreadsheet to a destination spreadsheet.

    All cell values and styles are copied over as well as row heights, column widths, merged cells.

    By default, the 1st worksheet is copied over.

    Args:
        src_path (Path): Source Excel spreadsheet path
        dest_path (Path): Destination Excel spreadsheet path
        source_sheet_index (int): Source spreadsheet worksheet index to copy to destination spreadsheet
    """
    from openpyxl.cell.cell import Cell

    src_book = openpyxl.load_workbook(src_path)
    dest_book = openpyxl.load_workbook(dest_path)
    sheet = src_book.worksheets[source_sheet_index]
    new_sheet = dest_book.create_sheet(sheet.title)
    dest_book.move_sheet(new_sheet, offset=(len(dest_book.sheetnames) * -2))
    for k, v in sheet.column_dimensions.items():
        new_sheet.column_dimensions[k] = copy(v)
    new_sheet.merged_cells = copy(sheet.merged_cells)
    for k, v in sheet.row_dimensions.items():
        new_sheet.row_dimensions[k] = copy(v)
    row: Tuple[Cell]
    for row in sheet.rows:
        for cell in row:
            new_cell: Cell = new_sheet[cell.coordinate]
            new_cell.value = cell.value
            if cell.has_style:
                new_cell.font = copy(cell.font)
                new_cell.border = copy(cell.border)
                new_cell.fill = copy(cell.fill)
                new_cell.number_format = copy(cell.number_format)
                new_cell.protection = copy(cell.protection)
                new_cell.alignment = copy(cell.alignment)
            new_cell.comment = copy(cell.comment)
    dest_book.save(filename=dest_path)


def write_xlsx_report(dfs: List[ExcelSheetDataFrame],
                      output_xlsx: Path,
                      quality_reqs: QualityRequirements,
                      images_for_sheets: Optional[List[SheetImage]] = None):
    """Write the output Excel XLSX file using the given dataframes.

    Args:
        dfs (List[ExcelSheetDataFrame]): List of ExcelSheetDataFrame objects to write to the output XLSX file
        output_xlsx (Path): Path to the output XLSX file
        quality_reqs (QualityRequirements): Quality requirements for the run
        images_for_sheets (Optional[List[SheetImage]]): List of SheetImage objects to add to the output XLSX file
    """
    with pd.ExcelWriter(output_xlsx, engine='xlsxwriter') as writer:
        monospace = dict(font_name='Courier New')
        text_wrap = dict(text_wrap=True)
        float_1dec = dict(num_format='0.0')
        perc_1dec = dict(num_format='0.0%')
        perc_2dec = dict(num_format='0.00%')

        book: Workbook = writer.book
        monospace_fmt = book.add_format(monospace)
        monospace_wrap_fmt = book.add_format({**monospace, **text_wrap})
        monospace_wrap_fmt.set_align('left')
        monospace_wrap_fmt.set_align('top')
        monospace_float_fmt = book.add_format({**monospace, **float_1dec})
        monospace_perc_fmt = book.add_format({**monospace, **perc_1dec})
        monospace_perc_2dec_fmt = book.add_format({**monospace, **perc_2dec})
        header_with_wrap_fmt = book.add_format({**text_wrap,
                                                'bold': True,
                                                'border': 1,
                                                'align': 'center',
                                                'valign': 'top'})
        varmap_header_fmt = book.add_format(dict(border=1,
                                                 align='left',
                                                 valign='bottom',
                                                 rotation=45,
                                                 font_name='Courier New'))
        fail_qc_fmt = book.add_format(dict(bg_color='FC9295',
                                           font_name='Courier New',
                                           bold=True))
        pass_qc_fmt = book.add_format(dict(bg_color='c4edce',
                                           font_name='Courier New',
                                           bold=False))
        float_cols = {'Mean Coverage Depth', 'Mean Depth'}
        perc_cols = {'% Genome Coverage'}
        perc_2dec_cols = {'Alternate Allele Frequency', 'Min AF', 'Max AF', 'Mean AF'}

        images_added = False

        for esdf in dfs:
            esdf.df = esdf.df.fillna('')  # Fill NaNs with empty strings to avoid issues with xlsxwriter when writing to Excel
            if images_for_sheets and esdf.sheet_name == SheetName.workflow_info.value:
                add_images(images_for_sheets, book)
                images_added = True
            esdf.df.to_excel(writer, sheet_name=esdf.sheet_name, **esdf.pd_to_excel_kwargs)

            sheet: Worksheet = book.get_worksheet_by_name(esdf.sheet_name)

            if esdf.pd_to_excel_kwargs.get('index', True):
                idx_and_cols = [esdf.df.index.name if esdf.df.index.name else ""] + list(esdf.df.columns)
            else:
                idx_and_cols = list(esdf.df.columns)

            for i, col_name in enumerate(idx_and_cols):
                if esdf.sheet_name != SheetName.consensus.value:
                    sheet.write_string(0, i, str(col_name), header_with_wrap_fmt)  # Force Bold Header for ALL sheets (Fixes Python 3.11+ issue), except consensus sheet.
                if esdf.header_comments and col_name in esdf.header_comments:
                    sheet.write_comment(0, i, esdf.header_comments[col_name])

            if esdf.autofit:
                for i, (width, col_name) in enumerate(zip(get_col_widths(esdf.df,
                                                                         index=True,
                                                                         max_width=80,
                                                                         include_header=esdf.include_header_width),
                                                          idx_and_cols)):
                    if col_name in float_cols:
                        logger.debug(f'{esdf.sheet_name}|Column {col_name} ({i}) width = {width}')
                        sheet.set_column(i, i, width, monospace_float_fmt)
                    elif col_name in perc_cols:
                        logger.debug(f'{esdf.sheet_name}|Column {col_name} ({i}) width = {width}')
                        sheet.set_column(i, i, width, monospace_perc_fmt)
                    elif col_name in perc_2dec_cols:
                        logger.debug(f'{esdf.sheet_name}|Column {col_name} ({i}) width = {width}')
                        sheet.set_column(i, i, width, monospace_perc_2dec_fmt)
                    else:
                        logger.debug(f'{esdf.sheet_name}|Column {col_name} ({i}) width = {width}')
                        sheet.set_column(i, i, width, monospace_fmt)
                if not esdf.include_header_width:
                    for i, col_name in enumerate(idx_and_cols):
                        sheet.write_string(0, i, string=col_name, cell_format=header_with_wrap_fmt)

            elif esdf.column_widths:
                for i, (width, col_name) in enumerate(zip(esdf.column_widths, idx_and_cols)):
                    logger.debug(f'{esdf.sheet_name}|Column {col_name} ({i}) width = {width}')
                    sheet.set_column(i, i, width, monospace_wrap_fmt)

                for i, idx in enumerate(esdf.df.index, start=1):
                    sheet.set_row(i, get_row_heights(esdf.df, idx), monospace_wrap_fmt)

            if esdf.sheet_name == SheetName.varmat.value:
                sheet.write_comment(row=0,
                                    col=0,
                                    comment=f'This sheet contains a matrix of alternate allele variant observation'
                                            f' frequency values for samples and variants. '
                                            f'3-colour conditional formatting is applied to the variant '
                                            f'frequency values where a major variant '
                                            f'(e.g. alternate allele frequency >={quality_reqs.major_allele_freq}) '
                                            f'is highlighted in green. Red indicates where the allele variant is not '
                                            f'observed in the sample (e.g. alternate allele frequency equals 0.0).')
                sheet.set_row(0, max(len(x) for x in idx_and_cols) * 5)
                for i, col_name in enumerate(idx_and_cols):
                    if i == 0:
                        continue
                    sheet.write_string(0, i, string=col_name, cell_format=varmap_header_fmt)
                sheet.conditional_format(first_row=1,
                                         first_col=1,
                                         last_row=esdf.df.shape[0],
                                         last_col=esdf.df.shape[1],
                                         options=dict(type='3_color_scale',
                                                      min_type='num',
                                                      mid_type='num',
                                                      max_type='num',
                                                      min_value=0.0,
                                                      mid_value=quality_reqs.major_allele_freq,
                                                      max_value=1.0))

            if esdf.sheet_name == SheetName.consensus.value:
                sheet.hide_gridlines(2)
                sheet.hide_row_col_headers()

            if esdf.sheet_name == SheetName.qc_stats.value:
                columns = esdf.df.columns.tolist()
                qc_status_column = 'QC Status'
                column_idx = columns.index(qc_status_column) + 1
                logger.info(f'Setting red background color for "{qc_status_column}" column (index={column_idx}) '
                            f'of sheet "{esdf.sheet_name}"')
                sheet.conditional_format(1, column_idx, esdf.df.shape[0], column_idx, options=dict(type='cell',
                                                                                                   value='"FAIL"',
                                                                                                   criteria='equal to',
                                                                                                   format=fail_qc_fmt))
                sheet.conditional_format(1, column_idx, esdf.df.shape[0], column_idx, options=dict(type='cell',
                                                                                                   value='"PASS"',
                                                                                                   criteria='equal to',
                                                                                                   format=pass_qc_fmt))
                cond_fmt_3color = dict(type='3_color_scale',
                                       min_type='num',
                                       mid_type='num',
                                       max_type='num', )

                add_cond_fmt(sheet,
                             esdf.df,
                             '% Genome Coverage',
                             {**cond_fmt_3color, **dict(min_value=0.0,
                                                        mid_value=quality_reqs.min_genome_coverage,
                                                        max_value=1.0)})
                for column in ['Median Coverage Depth', 'Mean Coverage Depth']:
                    add_cond_fmt(sheet,
                                 esdf.df,
                                 column,
                                 {**cond_fmt_3color, **dict(min_value=0,
                                                            mid_value=quality_reqs.min_median_depth,
                                                            max_type='max')}
                                 )
                ref_len = esdf.df['Reference Sequence Length'].values[0]
                for column in esdf.df.columns[esdf.df.columns.str.contains('X positions')]:
                    add_cond_fmt(sheet,
                                 esdf.df,
                                 column,
                                 {
                                     **cond_fmt_3color,
                                     **dict(min_color='6eb758',
                                            max_color='930e11',
                                            mid_color='f73639',
                                            min_value=0,
                                            mid_value=ref_len * 0.05,
                                            max_value=ref_len)}
                                 )
                for column in esdf.df.columns[esdf.df.columns.str.contains('Reads')]:
                    add_cond_fmt(sheet,
                                 esdf.df,
                                 column,
                                 {
                                     **cond_fmt_3color,
                                     **dict(min_color='930e11',
                                            mid_color='f4cf46',
                                            max_color='6eb758',
                                            min_value=0,
                                            mid_value=1000,
                                            max_type='max')}
                                 )

        if images_for_sheets is not None and not images_added:
            add_images(images_for_sheets, book)

    df_qc = get_qc_df(dfs)
    failed_samples = set(df_qc[df_qc['QC Status'] == 'FAIL'].index)
    add_comments(xlsx_path=output_xlsx, failed_samples=failed_samples, esdfs=dfs)


def add_cond_fmt(sheet: Worksheet,
                 df: pd.DataFrame,
                 column: str,
                 options: dict):
    columns = df.columns.tolist()
    column_idx = columns.index(column) + 1
    logger.info(f'Applying conditional formatting to column "{column}" (index={column_idx}) of '
                f'sheet "{sheet.name}"')

    sheet.conditional_format(first_row=1,
                             first_col=column_idx,
                             last_row=df.shape[0],
                             last_col=column_idx,
                             options=options)


def get_qc_df(dfs: List[ExcelSheetDataFrame]) -> Optional[pd.DataFrame]:
    for esdf in dfs:
        if esdf.sheet_name == SheetName.qc_stats.value:
            return esdf.df


def add_images(images_for_sheets: List[SheetImage],
               book: Workbook):
    """Add images and their descriptions to new sheets in a workbook"""
    import imageio
    text_wrap_fmt = book.add_format(dict(text_wrap=True, valign='justify'))
    for sheet_image in images_for_sheets:
        sheet = book.add_worksheet(sheet_image.sheet_name)
        sheet.set_column(0, 0, 100, text_wrap_fmt)
        sheet.write(0, 0, sheet_image.image_description, text_wrap_fmt)
        img = imageio.imread(sheet_image.image_path)
        x_size, y_size, _ = img.shape
        yx_ratio = y_size / x_size
        logger.debug(f'Image "{sheet_image.image_path.name}", x={x_size}, y={y_size}, y/x={yx_ratio}')
        sheet.insert_image(1, 0, sheet_image.image_path, options=dict(x_scale=1.0,
                                                                      y_scale=yx_ratio,
                                                                      object_position=3))
        sheet.hide_gridlines(2)
        sheet.hide_row_col_headers()


def get_excel_sheet_df(esds: List[ExcelSheetDataFrame],
                       sheet_name: str) -> Optional[ExcelSheetDataFrame]:
    for esd in esds:
        if esd.sheet_name == sheet_name:
            return esd


def add_comments(xlsx_path: Path,
                 failed_samples: Set[str],
                 esdfs: List[ExcelSheetDataFrame]) -> None:
    from openpyxl.comments import Comment
    from openpyxl.styles import PatternFill, Font
    from openpyxl.worksheet.worksheet import Worksheet
    from openpyxl.worksheet.dimensions import ColumnDimension
    logger.info(f'Loading workbook "{xlsx_path.name}" with openpyxl '
                f'to highlight {len(failed_samples)} samples that have failed QC')
    book = openpyxl.load_workbook(xlsx_path)
    logger.info(f'Loaded "{xlsx_path.name}" using openpyxl. Sheets: {book.sheetnames}')
    logger.info('Adjusting comment textbox sizes to fit text')
    for sheetname in book.sheetnames:
        for row in book[sheetname]:
            for cell in row:
                if cell.comment:
                    comment: Comment = cell.comment
                    comment.width = 300
                    comment.height = max((100, len(comment.text) / 3 * 2))
                    comment.author = f'xlavir version {__version__}'

    sheet_names = [
        SheetName.pangolin.value,
        SheetName.variants.value,
        SheetName.varmat.value,
    ]
    light_red = 'FC9295'
    for sheet_name in sheet_names:
        with contextlib.suppress(KeyError):
            sheet: Worksheet = book[sheet_name]
            logger.info(f'Highlighting failed samples in sheet "{sheet_name}".')
            for i, row in enumerate(sheet.rows):
                if i == 0:
                    continue
                cell = row[0]
                if cell.value in failed_samples:
                    cell.comment = Comment(f'Warning: Sample "{cell.value}" has failed general NGS QC',
                                           author='xlavir')
                    cell.fill = PatternFill(fill_type='solid',
                                            fgColor=light_red)
    esd_varmat = get_excel_sheet_df(esdfs, SheetName.varmat.value)
    esd_variants = get_excel_sheet_df(esdfs, SheetName.variants.value)
    if esd_varmat and esd_variants:
        with contextlib.suppress(KeyError):
            sheet: Worksheet = book[SheetName.varmat.value]
            logger.info('Adding additional comments to variant matrix values')
            df_varmat = esd_varmat.df
            variants: Dict[Tuple[str, str], Dict[str, Union[str, float, int]]] = esd_variants \
                .df.reset_index() \
                .set_index(['Sample', 'Mutation']) \
                .to_dict(orient='index')

            for i, row in enumerate(sheet.rows):
                if i == 0:
                    continue
                sample = df_varmat.index.values[i - 1]
                for j, cell in enumerate(row):
                    if j == 0:
                        continue
                    mutation = df_varmat.columns.values[j - 1]
                    if variant := variants.get((sample, mutation), None):
                        variant_str = '\n'.join(f'{k}: {v}' for k, v in variant.items())
                        comment_text = f'Sample: {sample}\nMutation: {mutation}\n{variant_str}'
                    else:
                        comment_text = f'Mutation "{mutation}" not found in sample "{sample}"'
                    cell.comment = Comment(comment_text,
                                           author=f'xlavir version {__version__}',
                                           width=300,
                                           height=len(comment_text))
    with contextlib.suppress(KeyError):
        sheet: Worksheet = book[SheetName.consensus.value]

        sheet.column_dimensions['A'] = ColumnDimension(worksheet=sheet,
                                                       index='A',
                                                       width=100)

        logger.info(f'Highlighting consensus sequences of failed '
                    f'samples in sheet "{SheetName.consensus.value}".')
        highlight_seq = False
        sample_name = ''
        dark_red = '260000'
        for row in sheet.rows:
            cell = row[0]
            font: Font = cell.font
            if cell.value[0] == '>':
                sample_name = cell.value[1:]
                if sample_name in failed_samples:
                    highlight_seq = True
                    # only add comment to cell containing failing sample fasta header
                    cell.comment = Comment(f'Warning: Sample "{sample_name}" has failed general NGS QC',
                                           author='xlavir')
                    cell.fill = PatternFill(fill_type='solid', fgColor=light_red)
                    cell.font = Font(name='Courier New',
                                     color=dark_red,
                                     size=font.size,
                                     family=font.family)
                else:
                    highlight_seq = False
                    cell.font = Font(name='Courier New',
                                     color='000000',
                                     size=font.size,
                                     family=font.family)
            elif cell.value and highlight_seq:
                cell.fill = PatternFill(fill_type='solid', fgColor=light_red)
                cell.font = Font(name='Courier New',
                                 color=dark_red,
                                 size=font.size,
                                 family=font.family)
            else:
                cell.font = Font(name='Courier New',
                                 color='000000',
                                 size=font.size,
                                 family=font.family)
    book.save(xlsx_path)
