import logging
import string
from typing import *
from collections import Counter

import numpy as np
import openpyxl
from openpyxl.styles import Alignment, Font, NamedStyle, PatternFill
from openpyxl.worksheet.worksheet import Worksheet
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA

from .. import QGRAIN_VERSION
from ..models import DistributionType, Dataset, Sample, ArtificialDataset, SSUResult, EMMAResult , UDMResult
from ..statistics import _all_scales, all_statistics, logarithmic
from ..distributions import get_distribution


def column_to_char(column_index: int):
    column = column_index + 1
    column_str = str()
    while column != 0:
        res = column % 26
        if res == 0:
            res = 26
            column -= 26
        column_str = chr(ord('A') + res - 1) + column_str
        column = column // 26
    return column_str


def to_cell_name(row: int, column: int):
    return "{0}{1}".format(column_to_char(column), row+1)


def prepare_styles(wb: openpyxl.Workbook):
    normal_light_style = NamedStyle(name="normal_light")
    normal_light_style.font = Font(size=12, name="Arial", color="000000")
    normal_light_style.alignment = Alignment(
        horizontal='center',
        vertical='center',
        text_rotation=0,
        wrap_text=False,
        shrink_to_fit=False,
        indent=0)
    normal_light_style.fill = PatternFill(
        patternType="solid",
        fgColor="c4cbcf")
    wb.add_named_style(normal_light_style)

    normal_dark_style = NamedStyle(name="normal_dark")
    normal_dark_style.font = Font(size=12, name="Arial", color="000000")
    normal_dark_style.alignment = Alignment(
        horizontal='center',
        vertical='center',
        text_rotation=0,
        wrap_text=False,
        shrink_to_fit=False,
        indent=0)
    normal_dark_style.fill = PatternFill(
        patternType="solid",
        fgColor="b2bbbe")
    wb.add_named_style(normal_dark_style)

    warning_style = NamedStyle(name="warning")
    warning_style.font = Font(size=12, name="Arial", color="ffffff")
    warning_style.alignment = Alignment(
        horizontal='center',
        vertical='center',
        text_rotation=0,
        wrap_text=False,
        shrink_to_fit=False,
        indent=0)
    warning_style.fill = PatternFill(
        patternType="solid",
        fgColor="ff9900")
    wb.add_named_style(warning_style)

    header_style = NamedStyle(name="header")
    header_style.font = Font(size=14, bold=True, name="Times New Roman", color="ffffff")
    header_style.alignment = Alignment(
        horizontal='center',
        vertical='center',
        text_rotation=0,
        wrap_text=True,
        shrink_to_fit=False,
        indent=0)
    header_style.fill = PatternFill(
        patternType="solid",
        fgColor="475164")
    wb.add_named_style(header_style)

    description_style = NamedStyle(name="description")
    description_style.font = Font(size=14, bold=True, name="Times New Roman", color="ffffff")
    description_style.alignment = Alignment(
        horizontal='left',
        vertical='center',
        text_rotation=0,
        wrap_text=False,
        shrink_to_fit=False,
        indent=0)
    description_style.fill = PatternFill(
        patternType="solid",
        fgColor="475164")
    wb.add_named_style(description_style)


SMALL_WIDTH = 12
MEDIAN_WIDTH = 24
LARGE_WIDTH = 48


def _write_readme_sheet(ws: Worksheet, text: str):
    ws.title = "README"
    full_text = \
        """
        This Excel file was generated by QGrain ({0}).

        Please cite:
        Liu, Y., Liu, X., Sun, Y., 2021. QGrain: An open-source and easy-to-use software for the comprehensive
            analysis of grain size distributions. Sedimentary Geology 423, 105980. DOI: 10.1016/j.sedgeo.2021.105980

        {1}

        """.format(QGRAIN_VERSION, text)
    lines_of_description = full_text.split("\n")
    for row, line in enumerate(lines_of_description):
        cell = ws.cell(row+1, 1, value=line)
        cell.style = "description"
    ws.column_dimensions[column_to_char(0)].width = 200


def _write_dataset_sheet(ws: Worksheet, dataset: Dataset, progress_callback: Callable[[float], None] = None):
    def write(r, c, v, s="normal_light"):
        cell = ws.cell(r+1, c+1, value=v)
        cell.style = s

    ws.title = "GSDs"
    write(0, 0, "Sample Name", s="header")
    ws.column_dimensions[column_to_char(0)].width = 24
    for col, value in enumerate(dataset.classes, 1):
        write(0, col, value, s="header")
        ws.column_dimensions[column_to_char(col)].width = 10

    for i_sample, sample in enumerate(dataset):
        row = i_sample + 1
        if row % 2 == 0:
            style = "normal_dark"
        else:
            style = "normal_light"
        write(row, 0, sample.name, s=style)
        for col, value in enumerate(sample.distribution, 1):
            write(row, col, value, s=style)
        if progress_callback is not None:
            progress = i_sample / len(dataset)
            progress_callback(progress)


def save_artificial_dataset(dataset: ArtificialDataset, filename: str,
                            progress_callback: Callable[[float], None] = None,
                            logger: logging.Logger = None):
    if logger is None:
        logger = logging.getLogger("QGrain")
    else:
        assert isinstance(logger, logging.Logger)
    logger.debug("Start to save the artificial dataset.")
    wb = openpyxl.Workbook()
    prepare_styles(wb)
    logger.debug("Creating the README sheet.")
    readme_text = \
        f"""
        It contains {dataset.n_components+2} sheets:
        1. The first sheet is the grain size distributions of this artificial dataset.
        2. The second sheet stores the parameters which were used to generate the dataset.
        3-{dataset.n_components+2}. Each of the left sheets is used to put the distributions of each component.

        Sampling Settings
            Minimum size [micron]: {dataset.classes[0]}
            Maximum size [micron]: {dataset.classes[-1]}
            Number of Grain Size Classes: {len(dataset.classes)}
            Precision: {dataset.precision}
            Noise Decimals: {dataset.noise}
            Number of Samples: {len(dataset)}
        """
    _write_readme_sheet(wb.active, readme_text)

    def write(r, c, v, s="normal_light"):
        cell = ws.cell(r + 1, c + 1, value=v)
        cell.style = s

    distribution_class = get_distribution(dataset.distribution_type)
    parameter_names = list(distribution_class.PARAMETER_NAMES) + ["Weight"]

    logger.debug("Creating the GSDs sheet.")
    ws = wb.create_sheet("GSDs")
    write(0, 0, "Sample Name", s="header")
    ws.column_dimensions[column_to_char(0)].width = 24
    for col, value in enumerate(dataset.classes, 1):
        write(0, col, value, s="header")
        ws.column_dimensions[column_to_char(col)].width = 10
    for row, sample in enumerate(dataset, 1):
        if row % 2 == 0:
            style = "normal_dark"
        else:
            style = "normal_light"
        write(row, 0, sample.name, s=style)
        for col, value in enumerate(sample.distribution, 1):
            write(row, col, value, s=style)
        if progress_callback is not None:
            progress_callback(row / len(dataset) * 0.2)

    logger.debug("Creating the Parameters sheet.")
    ws = wb.create_sheet("Parameters")
    write(0, 0, "Sample Name", s="header")
    ws.merge_cells(start_row=1, start_column=1, end_row=2, end_column=1)
    ws.column_dimensions[column_to_char(0)].width = 24
    for i in range(dataset.n_components):
        write(0, dataset.n_parameters * i + 1, f"Component{i+1}", s="header")
        ws.merge_cells(start_row=1, start_column=dataset.n_parameters*i+2,
                       end_row=1, end_column=dataset.n_parameters*(i+1)+1)
        for j, header_name in enumerate(parameter_names):
            write(1, dataset.n_parameters * i + 1 + j, header_name, s="header")
            ws.column_dimensions[column_to_char(dataset.n_parameters*i+1+j)].width = 16
    for i, sample in enumerate(dataset):
        row = i + 2
        if row % 2 == 1:
            style = "normal_dark"
        else:
            style = "normal_light"
        write(row, 0, sample.name, s=style)
        for j, component in enumerate(sample):
            for k in range(dataset.n_parameters):
                write(row, dataset.n_parameters * j + k + 1, dataset.parameters[i, k, j], s=style)
        if progress_callback is not None:
            progress_callback(1 / len(dataset) * 0.2 + 0.2)

    for i in range(dataset.n_components):
        logger.debug(f"Creating the C{i+1} sheet.")
        ws = wb.create_sheet(f"C{i+1}")
        write(0, 0, "Sample Name", s="header")
        ws.column_dimensions[column_to_char(0)].width = 24
        for col, value in enumerate(dataset.classes, 1):
            write(0, col, value, s="header")
            ws.column_dimensions[column_to_char(col)].width = 10
        for row, sample in enumerate(dataset, 1):
            if row % 2 == 0:
                style = "normal_dark"
            else:
                style = "normal_light"
            write(row, 0, sample.name, s=style)
            for col, value in enumerate(sample[i].distribution, 1):
                write(row, col, value, s=style)
            if progress_callback is not None:
                progress_callback(((i*len(dataset) + row) / len(dataset)*dataset.n_components) * 0.6 + 0.4)

    wb.save(filename)
    wb.close()
    if progress_callback is not None:
        progress_callback(1.0)
    logger.info(f"The artificial dataset has been saved to the Excel file: [{filename}].")


def save_dataset(dataset: Dataset, filename: str, progress_callback: Callable[[float], None] = None,
                 logger: logging.Logger = None):
    assert dataset is not None
    if logger is None:
        logger = logging.getLogger("QGrain")
    else:
        assert isinstance(logger, logging.Logger)
    logger.debug("Start to save the grain size dataset.")
    wb = openpyxl.Workbook()
    prepare_styles(wb)
    logger.debug("Creating the README sheet.")
    readme_text = \
        """
        It only contains one sheet which stores the grain size distributions.
        """
    _write_readme_sheet(wb.active, readme_text)
    logger.debug("Creating the GSDs sheet.")
    ws = wb.create_sheet("GSDs")
    _write_dataset_sheet(ws, dataset, progress_callback)
    logger.debug("Saving the workbook to file.")
    wb.save(filename)
    wb.close()
    if progress_callback is not None:
        progress_callback(1.0)
    logger.info(f"The dataset has been saved to the Excel file: [{filename}].")


def save_statistics(dataset: Dataset, filename: str, progress_callback: Callable[[float], None] = None,
                    logger: logging.Logger = None):
    assert dataset is not None
    if logger is None:
        logger = logging.getLogger("QGrain")
    else:
        assert isinstance(logger, logging.Logger)
    logger.debug("Start to save the statistical result.")
    # Calculate
    logger.debug("Calculating the statistical parameters and classification groups of all samples.")
    all_sample_statistics = []
    for i, sample in enumerate(dataset):
        sample_statistics = all_statistics(sample.classes, sample.classes_phi, sample.distribution)
        all_sample_statistics.append(sample_statistics)
        if progress_callback is not None:
            progress = (i / len(dataset)) * 0.4
            progress_callback(progress)
    wb = openpyxl.Workbook()
    prepare_styles(wb)
    logger.debug("Creating the README sheet.")
    readme_text = \
        """
        It contains 6 sheets.
            1-5. The previous five sheets stores the statistical parameters of five computational methods, respectively.
            6. The last sheet puts the proportions of different size scales and the classification groups.

        The statistical formulas are referred to Blott & Pye (2001)'s work.
        The classification of GSDs is referred to Folk (1954)'s and Blott & Pye (2012)'s scheme.

        References:
            1. Blott, S. J. & Pye, K. Particle size scales and classification of sediment types based on particle size
                distributions: Review and recommended procedures. Sedimentology 59, 2071–2096 (2012).
            2. Blott, S. J. & Pye, K. GRADISTAT: a grain-size distribution and statistics package for the analysis of
                unconsolidated sediments. Earth Surf. Process. Landforms 26, 1237–1248 (2001).
            3. Folk, R. L. The Distinction between Grain Size and Mineral Composition in Sedimentary-Rock Nomenclature.
                The Journal of Geology 62, 344–359 (1954).
        """
    _write_readme_sheet(wb.active, readme_text)
    logger.debug("Creating the GSDs sheet.")
    if progress_callback is not None:
        _callback = lambda p: progress_callback(p * 0.1 + 0.4)
    else:
        _callback = None
    ws = wb.create_sheet("GSDs")
    _write_dataset_sheet(ws, dataset, progress_callback=_callback)

    def get_keys(method: str):
        small_width = 12
        median_width = 24
        large_width = 48
        if method == "arithmetic":
            unit = "μm"
            keys = [
                (lambda s: s[method]["mean"], f"Mean [{unit}]", small_width),
                (lambda s: s[method]["std"], f"Sorting Coefficient", small_width),
                (lambda s: s[method]["skewness"], f"Skewness", small_width),
                (lambda s: s[method]["kurtosis"], f"Kurtosis", small_width)]
            return keys
        elif method in ("geometric", "logarithmic", "geometric_fw57", "logarithmic_fw57"):
            unit = "μm" if method.startswith("geometric") else "φ"
            keys = [
                (lambda s: s[method]["mean"], f"Mean [{unit}]", small_width),
                (lambda s: s[method]["mean_description"], "Mean Description", median_width),
                (lambda s: s[method]["first_percentile"], f"First Percentile (C) [{unit}]", small_width),
                (lambda s: s[method]["median"], f"Median (M) [{unit}]", small_width),
                (lambda s: s[method]["mode"], f"Mode [{unit}]", small_width),
                (lambda s: len(s[method]["modes"]), f"Number of Modes", small_width),
                (lambda s: ", ".join([f"{mode: 0.4f}" for mode in s[method]["modes"]]), f"Modes [{unit}]", median_width),
                (lambda s: s[method]["proportion_ss"] * 100.0, "Sortable Silt Proportion [%]", small_width),
                (lambda s: s[method]["mean_ss"], f"Mean Sortable Silt Size [{unit}]", small_width),
                (lambda s: s[method]["std"], "Sorting Coefficient", small_width),
                (lambda s: s[method]["std_description"], "Sorting Description", median_width),
                (lambda s: s[method]["skewness"], "Skewness", small_width),
                (lambda s: s[method]["skewness_description"], "Skewness Description", median_width),
                (lambda s: s[method]["kurtosis"], "Kurtosis", small_width),
                (lambda s: s[method]["kurtosis_description"], "Kurtosis Description", median_width)]
            return keys
        elif method == "proportion_and_classification":
            keys = [
                (lambda s: ", ".join([f"{p*100:0.4f}" for p in s["proportions_gsm"]]), "(Gravel, Sand, Mud) Proportions [%]", large_width),
                (lambda s: ", ".join([f"{p*100:0.4f}" for p in s["proportions_ssc"]]), "(Sand, Silt, Clay) Proportions [%]", large_width),
                (lambda s: ", ".join([f"{p*100:0.4f}" for p in s["proportions_bgssc"]]), "(Boulder, Gravel, Sand, Silt, Clay) Proportions [%]", large_width),
                (lambda s: s["group_folk54"], "Group (Folk, 1954)", median_width),
                (lambda s: s["group_bp12"], "Group (Blott & Pye, 2012)", large_width),
                (lambda s: s["group_bp12_symbol"], "Group Symbol (Blott & Pye, 2012)", median_width)]
            all_scales = _all_scales()
            for scale in all_scales:
                func = lambda s, scale=scale: s["proportions"][scale] * 100.0
                name = string.capwords(" ".join(scale)) + " Proportion [%]"
                keys.append((func, name, small_width))
            return keys
        else:
            raise NotImplementedError(method)

    def write(r, c, v, s="normal_light"):
        cell = ws.cell(r + 1, c + 1, value=v)
        cell.style = s

    methods = ["arithmetic", "geometric", "logarithmic", "geometric_fw57",
               "logarithmic_fw57", "proportion_and_classification"]
    sheet_names = ["Arithmetic", "Geometric", "Logarithmic", "Geometric_fw57",
                   "Logarithmic_fw57", "Proportion and Classification"]
    for i_method, (method, sheet_name) in enumerate(zip(methods, sheet_names)):
        logger.debug(f"Creating the {sheet_name} sheet.")
        ws = wb.create_sheet(sheet_name)
        write(0, 0, "Sample Name", s="header")
        keys = get_keys(method)
        ws.column_dimensions[column_to_char(0)].width = 16
        for col, (func, name, width) in enumerate(keys, 1):
            write(0, col, name, s="header")
            ws.column_dimensions[column_to_char(col)].width = width
        for i_sample, (sample, sample_statistics) in enumerate(zip(dataset, all_sample_statistics)):
            row = i_sample + 1
            if row % 2 == 0:
                style = "normal_dark"
            else:
                style = "normal_light"
            write(row, 0, sample.name, s=style)
            for col, (func, name, width) in enumerate(keys, 1):
                value = func(sample_statistics)
                write(row, col, value, s=style)
            if progress_callback is not None:
                progress = 0.5 + ((i_sample / len(dataset)) + i_method) / len(methods) * 0.5
                progress_callback(progress)
    logger.debug("Saving the workbook to file.")
    wb.save(filename)
    wb.close()
    if progress_callback is not None:
        progress_callback(1.0)
    logger.info(f"The statistical result has been saved to the Excel file: [{filename}].")


def save_pca(dataset: Dataset, filename: str, progress_callback: Callable[[float], None] = None,
             logger: logging.Logger = None):
    assert dataset is not None
    if logger is None:
        logger = logging.getLogger("QGrain")
    else:
        assert isinstance(logger, logging.Logger)
    logger.debug("Start to save the PCA result.")
    # Calculate
    logger.debug("Performing the PCA algorithm to this dataset. The number of PCs is set to 10.")
    pca = PCA(n_components=10)
    transformed = pca.fit_transform(dataset.distributions)
    components = pca.components_
    ratios = pca.explained_variance_ratio_

    wb = openpyxl.Workbook()
    prepare_styles(wb)
    logger.debug("Creating the README sheet.")
    readme_text = \
        """
        It contains 3 sheets:
        1. The first sheet is the dataset which was used to perform the PCA algorithm.
        2. The second sheet is used to put the distributions of all PCs.
        3. The third sheet is used to store the PC variations of all samples.

        The base PCA algorithm is implemented by scikit-learn.
        You can get the details of algorithm from the following website.
        https://scikit-learn.org/stable/modules/generated/sklearn.decomposition.PCA.html
        """
    _write_readme_sheet(wb.active, readme_text)
    logger.debug("Creating the GSDs sheet.")
    if progress_callback is not None:
        _callback = lambda progress: progress_callback(progress * 0.2)
    else:
        _callback = None
    ws = wb.create_sheet("GSDs")
    _write_dataset_sheet(ws, dataset, progress_callback=_callback)

    def write(r, c, v, s="normal_light"):
        cell = ws.cell(r + 1, c + 1, value=v)
        cell.style = s

    logger.debug("Creating the Distributions of PCs sheet.")
    ws = wb.create_sheet("Distributions of PCs")
    write(0, 0, "PC", s="header")
    ws.column_dimensions[column_to_char(0)].width = 16
    for col, value in enumerate(dataset.classes, 1):
        write(0, col, value, s="header")
        ws.column_dimensions[column_to_char(col)].width = 10
    for i, component in enumerate(components):
        row = i + 1
        if row % 2 == 0:
            style = "normal_dark"
        else:
            style = "normal_light"
        write(row, 0, f"PC{i+1} ({ratios[i]:0.2%})", s=style)
        for col, value in enumerate(component, 1):
            write(row, col, value, s=style)
    if progress_callback is not None:
        progress_callback(0.3)

    logger.debug("Creating the Variations of PCs sheet.")
    ws = wb.create_sheet("Variations of PCs")
    write(0, 0, "Sample Name", s="header")
    ws.column_dimensions[column_to_char(0)].width = 16
    for i in range(10):
        write(0, i + 1, f"PC{i+1} ({ratios[i]:0.2%})", s="header")
        ws.column_dimensions[column_to_char(i+1)].width = 10
    for row, variations in enumerate(transformed, 1):
        if row % 2 == 0:
            style = "normal_dark"
        else:
            style = "normal_light"
        write(row, 0, dataset[row-1].name, s=style)
        for col, value in enumerate(variations, 1):
            write(row, col, value, s=style)
        if progress_callback is not None:
            progress_callback(row / len(dataset) * 0.7 + 0.3)

    wb.save(filename)
    wb.close()
    if progress_callback is not None:
        progress_callback(1.0)
    logger.info(f"The PCA result has been saved to the Excel file: [{filename}].")


def save_clustering(dataset: Dataset, flags: Sequence[int], filename: str,
                    data_key: Tuple, linkage_method: str, metric: str,
                    progress_callback: Callable[[float], None] = None, logger: logging.Logger = None):
    assert dataset is not None
    if logger is None:
        logger = logging.getLogger("QGrain")
    else:
        assert isinstance(logger, logging.Logger)
    logger.debug("Start to save the clustering result.")
    flag_set = set(flags)
    n_clusters = len(flag_set)
    typical_samples: List[Sample] = []
    temp_flag_set = set()
    for i, flag in enumerate(flags):
        if len(temp_flag_set) == n_clusters:
            break
        if flag not in temp_flag_set:
            typical_samples.append(dataset[i])
            temp_flag_set.add(flag)

    wb = openpyxl.Workbook()
    prepare_styles(wb)
    logger.debug("Creating the README sheet.")
    readme_text = \
        f"""
        It contains 3 or {n_clusters+3} sheets:
        1. The first sheet is the dataset which was used to perform the hierarchy clustering algorithm.
        2. The second sheet is used to put the clustering flags of all samples.
        3. The third sheet is the typical samples (i.e, the first sample of each cluster was selected).
        4. If the number of clusters less equal to 100, the samples of each cluster will be save to individual sheets.

        The base hierarchy clustering algorithm is implemented by Scipy.
        You can get the details of algorithm from the following website.
        https://docs.scipy.org/doc/scipy/reference/cluster.hierarchy.html
        
        Data key: {data_key}
        Linkage method: {linkage_method}
        Metric: {metric}
        """
    _write_readme_sheet(wb.active, readme_text)
    logger.debug("Creating the GSDs sheet.")
    if progress_callback is not None:
        _callback = lambda p: progress_callback(p * 0.2)
    else:
        _callback = None
    ws = wb.create_sheet("GSDs")
    _write_dataset_sheet(ws, dataset, progress_callback=_callback)

    def write(r, c, v, s="normal_light"):
        cell = ws.cell(r + 1, c + 1, value=v)
        cell.style = s

    logger.debug("Creating the Cluster Flags of Samples sheet.")
    ws = wb.create_sheet("Cluster Flags of Samples")
    write(0, 0, "Sample Name", s="header")
    write(0, 1, "Cluster Flags", s="header")
    ws.column_dimensions[column_to_char(0)].width = 16
    ws.column_dimensions[column_to_char(1)].width = 16
    for i, (sample, flag) in enumerate(zip(dataset, flags)):
        row = i + 1
        if row % 2 == 0:
            style = "normal_dark"
        else:
            style = "normal_light"
        write(row, 0, sample.name, s=style)
        write(row, 1, flag, s=style)
        if progress_callback is not None:
            if n_clusters <= 100:
                progress_callback(i / len(dataset) * 0.1 + 0.2)
            else:
                progress_callback(i / len(dataset) * 0.4 + 0.2)

    logger.debug("Creating the Typical Samples of Clusters sheet.")
    ws = wb.create_sheet("Typical Samples of Clusters")
    write(0, 0, "Sample Name", s="header")
    ws.column_dimensions[column_to_char(0)].width = 16
    for col, value in enumerate(dataset.classes, 1):
        write(0, col, value, s="header")
        ws.column_dimensions[column_to_char(col)].width = 10
    for i, sample in enumerate(typical_samples):
        row = i + 1
        if row % 2 == 0:
            style = "normal_dark"
        else:
            style = "normal_light"
        write(row, 0, sample.name, s=style)
        for col, value in enumerate(sample.distribution, 1):
            write(row, col, value, s=style)
        if progress_callback is not None:
            if n_clusters <= 100:
                progress_callback(i / len(dataset) * 0.1 + 0.3)
            else:
                progress_callback(i / len(dataset) * 0.4 + 0.6)

    if n_clusters <= 100:
        for flag in flag_set:
            samples = []
            for sample, in_this_cluster in zip(dataset, np.equal(flags, flag)):
                if in_this_cluster:
                    samples.append(sample)
            logger.debug(f"Creating the Cluster{flag} sheet.")
            ws = wb.create_sheet(f"Cluster{flag}")
            write(0, 0, "Sample Name", s="header")
            ws.column_dimensions[column_to_char(0)].width = 16
            for col, value in enumerate(dataset.classes, 1):
                write(0, col, value, s="header")
                ws.column_dimensions[column_to_char(col)].width = 10
            for i, sample in enumerate(samples):
                row = i + 1
                if row % 2 == 0:
                    style = "normal_dark"
                else:
                    style = "normal_light"
                write(row, 0, sample.name, s=style)
                for col, value in enumerate(sample.distribution, 1):
                    write(row, col, value, s=style)
                if progress_callback is not None:
                    progress_callback(i / len(dataset) / n_clusters * 0.6 + 0.4)

    wb.save(filename)
    wb.close()
    if progress_callback is not None:
        progress_callback(1.0)
    logger.info(f"The Clustering result has been saved to the Excel file: [{filename}].")


def save_emma(result: EMMAResult, filename: str, progress_callback: Callable[[float], None] = None,
              logger: logging.Logger = None):
    if logger is None:
        logger = logging.getLogger("QGrain")
    else:
        assert isinstance(logger, logging.Logger)
    logger.debug("Start to save the EMMA result.")

    wb = openpyxl.Workbook()
    prepare_styles(wb)
    logger.debug("Creating the README sheet.")
    readme_text = \
        f"""
        It contains three sheets:
        1. The first sheet is the dataset which was used to perform the EMMA algorithm.
        2. The second sheet is used to put the distributions of all end members.
        3. The third sheet is the proportions of end members of all samples.

        This EMMA algorithm was implemented by QGrain, using the famous machine learning framework, PyTorch.

        EMMA Algorithm Details
            Number of Samples: {result.n_samples}
            Kernel Type: {result.kernel_type.name}
            Number of End Members: {result.n_members}
            x0: {result.x0 if result.x0 is None else result.x0.tolist()}
            Number of Epochs: {len(result.loss_series("lmse"))}
            Spent Time: {result.time_spent} s

            Computing Device: {"Unknown" if result.settings is None else result.settings["device"]}
            Loss Function: {"Unknown" if result.settings is None else result.settings["loss"]}
            Number of Pretrain Epochs: {"Unknown" if result.settings is None else result.settings["pretrain_epochs"]}
            Minimum Number of Epochs: {"Unknown" if result.settings is None else result.settings["min_epochs"]}
            Maximum Number of Epochs: {"Unknown" if result.settings is None else result.settings["max_epochs"]}
            Learning Rate: {"Unknown" if result.settings is None else result.settings["learning_rate"]}
            Precision: {"Unknown" if result.settings is None else result.settings["precision"]}
            Betas: {"Unknown" if result.settings is None else result.settings["betas"]}
            Update End Members: {"Unknown" if result.settings is None else result.settings["update_end_members"]}
            Need History: {"Unknown" if result.settings is None else result.settings["need_history"]}
        """
    _write_readme_sheet(wb.active, readme_text)
    logger.debug("Creating the GSDs sheet.")
    if progress_callback is not None:
        _callback = lambda p: progress_callback(p * 0.4)
    else:
        _callback = None
    ws = wb.create_sheet("GSDs")
    _write_dataset_sheet(ws, result.dataset, progress_callback=_callback)

    def write(r, c, v, s="normal_light"):
        cell = ws.cell(r + 1, c + 1, value=v)
        cell.style = s

    logger.debug("Creating the Distributions of End Members sheet.")
    ws = wb.create_sheet("Distributions of End Members")
    write(0, 0, "End Member", s="header")
    ws.column_dimensions[column_to_char(0)].width = 16
    for col, value in enumerate(result.dataset.classes, 1):
        write(0, col, value, s="header")
        ws.column_dimensions[column_to_char(col)].width = 10
    for i in range(result.n_members):
        row = i + 1
        if row % 2 == 0:
            style = "normal_dark"
        else:
            style = "normal_light"
        write(row, 0, f"EM{i+1}", s=style)
        for col, value in enumerate(result.end_members[i], 1):
            write(row, col, value, s=style)

    logger.debug("Creating the Proportions of End Members sheet.")
    ws = wb.create_sheet("Proportions of End Members")
    write(0, 0, "Sample Name", s="header")
    ws.column_dimensions[column_to_char(0)].width = 16
    for i in range(result.n_members):
        write(0, i + 1, f"EM{i+1} [%]", s="header")
        ws.column_dimensions[column_to_char(i+1)].width = 10
    for i, sample_proportions in enumerate(result.proportions):
        row = i + 1
        if row % 2 == 0:
            style = "normal_dark"
        else:
            style = "normal_light"
        write(row, 0, result.dataset[i].name, s=style)
        for j in range(result.n_members):
            col = j + 1
            write(row, col, sample_proportions[j]*100, s=style)
        if progress_callback is not None:
            progress_callback(i / result.n_samples * 0.6 + 0.4)

    wb.save(filename)
    wb.close()
    if progress_callback is not None:
        progress_callback(1.0)
    logger.info(f"The EMMA result has been saved to the Excel file: [{filename}].")


def save_ssu(results: List[SSUResult], filename: str, align_components=False,
             progress_callback: Callable[[float], None] = None, logger: logging.Logger = None):
    if logger is None:
        logger = logging.getLogger("QGrain")
    else:
        assert isinstance(logger, logging.Logger)
    logger.debug("Start to save the SSU results.")
    # pack the GSDs of samples to the dataset
    dataset = Dataset(
        "GSDs", [result.sample.name for result in results],
        results[0].classes,
        [result.sample.distribution for result in results])
    max_n_components = max(Counter([len(result) for result in results]).keys())

    # prepare flags
    flags = []
    if not align_components:
        for result in results:
            flags.extend(range(len(result)))
    else:
        stacked_components = []
        for result in results:
            for component in result:
                stacked_components.append(component.distribution)
        stacked_components = np.array(stacked_components)
        clusters = KMeans(n_clusters=max_n_components)
        flags = clusters.fit_predict(stacked_components)
        # check flags to make it unique
        flag_index = 0
        for result in results:
            result_flags = set()
            for i in range(len(result)):
                if flags[flag_index] in result_flags:
                    if flags[flag_index] == max_n_components:
                        flags[flag_index] = max_n_components-1
                    else:
                        flags[flag_index] += 1
                    result_flags.add(flags[flag_index])
                flag_index += 1

        flag_set = set(flags)
        picked = []
        for target_flag in flag_set:
            for i, flag in enumerate(flags):
                if flag == target_flag:
                    picked.append((target_flag, logarithmic(dataset.classes_phi, stacked_components[i])["mean"]))
                    break
        picked.sort(key=lambda x: x[1])
        flag_map = {flag: index for index, (flag, _) in enumerate(picked)}
        flags = np.array([flag_map[flag] for flag in flags])

    wb = openpyxl.Workbook()
    prepare_styles(wb)
    logger.debug("Creating the README sheet.")
    readme_text = \
        f"""
        It contains {max_n_components+4} sheets:
        1. The first sheet is used to put the grain size distributions of corresponding samples.
        2. The second sheet is used to put the fitting information and resolved parameters of SSU results.
        3. The third sheet stores the statistical parameters of all components.
        4. The fourth sheet is used to put the distributions of unmixed components and the sum.
        5-{max_n_components+4}. Other sheets severally store the distributions of each component group.

        The SSU algorithm is implemented by QGrain.
        """
    _write_readme_sheet(wb.active, readme_text)
    logger.debug("Creating the GSDs sheet.")
    if progress_callback is not None:
        _callback = lambda p: progress_callback(p * 0.1)
    else:
        _callback = None
    ws = wb.create_sheet("GSDs")
    _write_dataset_sheet(ws, dataset, progress_callback=_callback)

    def write(r, c, v, s="normal_light"):
        cell = ws.cell(r + 1, c + 1, value=v)
        cell.style = s

    logger.debug("Creating the Information of Fitting sheet.")
    ws = wb.create_sheet("Information of Fitting")
    write(0, 0, "Sample Name", s="header")
    ws.column_dimensions[column_to_char(0)].width = 16
    headers = ["Distribution Type",
               "Number of Components",
               "Resolver Settings",
               "x0",
               "Resolved Parameters",
               "Spent Time [s]",
               "Number of Iterations",
               "Final Loss [LMSE]"]
    for col, value in enumerate(headers, 1):
        write(0, col, value, s="header")
        if col in (3, 4, 5):
            ws.column_dimensions[column_to_char(col)].width = 32
        else:
            ws.column_dimensions[column_to_char(col)].width = 10
    for i, result in enumerate(results):
        row = i + 1
        if row % 2 == 0:
            style = "normal_dark"
        else:
            style = "normal_light"
        write(row, 0, result.sample.name, s=style)
        write(row, 1, result.distribution_type.name, s=style)
        write(row, 2, len(result), s=style)
        write(row, 3, "Default" if result.settings is None else str(result.settings), s=style)
        write(row, 4, "None" if result.x0 is None else str(result.x0.tolist()), s=style)
        write(row, 5, str(result.parameters[-1].tolist()), s=style)
        write(row, 6, result.time_spent, s=style)
        write(row, 7, len(result.loss_series("lmse")), s=style)
        write(row, 8, result.loss("lmse"), s=style)
        if progress_callback is not None:
            progress_callback(i / len(dataset) * 0.1 + 0.1)

    logger.debug("Creating the Statistical Moments sheet.")
    ws = wb.create_sheet("Statistical Moments")
    write(0, 0, "Sample Name", s="header")
    ws.merge_cells(start_row=1, start_column=1, end_row=2, end_column=1)
    ws.column_dimensions[column_to_char(0)].width = 16
    headers = []
    sub_headers = ["Proportion [%]",
                   "Mean [phi]",
                   "Mean [microns]",
                   "Standard Deviation [phi]",
                   "Standard Deviation [microns]",
                   "Skewness",
                   "Kurtosis"]
    for i in range(max_n_components):
        write(0, i * len(sub_headers) + 1, f"C{i+1}", s="header")
        ws.merge_cells(start_row=1, start_column=i*len(sub_headers)+2, end_row=1, end_column=(i+1)*len(sub_headers)+1)
        headers.extend(sub_headers)
    for col, value in enumerate(headers, 1):
        write(1, col, value, s="header")
        ws.column_dimensions[column_to_char(col)].width = 10
    flag_index = 0
    for i, result in enumerate(results):
        row = i + 2
        if row % 2 == 0:
            style = "normal_light"
        else:
            style = "normal_dark"
        write(row, 0, result.sample.name, s=style)
        for component in result:
            index = flags[flag_index]
            s = all_statistics(result.classes, result.classes_phi, component.distribution)
            write(row, index * len(sub_headers) + 1, component.proportion*100, s=style)
            write(row, index * len(sub_headers) + 2, s["logarithmic"]["mean"], s=style)
            write(row, index * len(sub_headers) + 3, s["geometric"]["mean"], s=style)
            write(row, index * len(sub_headers) + 4, s["logarithmic"]["std"], s=style)
            write(row, index * len(sub_headers) + 5, s["geometric"]["std"], s=style)
            write(row, index * len(sub_headers) + 6, s["logarithmic"]["skewness"], s=style)
            write(row, index * len(sub_headers) + 7, s["logarithmic"]["kurtosis"], s=style)
            flag_index += 1
        if progress_callback is not None:
            progress_callback(i / len(dataset) * 0.1 + 0.2)

    logger.debug("Creating the Unmixed Components sheet.")
    ws = wb.create_sheet("Unmixed Components")
    ws.merge_cells(start_row=1, start_column=1, end_row=1, end_column=2)
    write(0, 0, "Sample Name", s="header")
    ws.column_dimensions[column_to_char(0)].width = 16
    for col, value in enumerate(dataset.classes, 2):
        write(0, col, value, s="header")
        ws.column_dimensions[column_to_char(col)].width = 10
    row = 1
    for i, result in enumerate(results):
        if i % 2 == 0:
            style = "normal_light"
        else:
            style = "normal_dark"
        write(row, 0, result.sample.name, s=style)
        ws.merge_cells(start_row=row+1, start_column=1, end_row=row+len(result)+1, end_column=1)
        for component_i, component in enumerate(result, 1):
            write(row, 1, f"C{component_i}", s=style)
            for col, value in enumerate(component.distribution*component.proportion, 2):
                write(row, col, value, s=style)
            row += 1
        write(row, 1, "Sum", s=style)
        for col, value in enumerate(result.distribution, 2):
            write(row, col, value, s=style)
        row += 1
        if progress_callback is not None:
            progress_callback(i / len(dataset) * 0.2 + 0.3)

    logger.debug("Creating separate sheets for all components.")
    ws_dict = {}
    flag_set = set(flags)
    for flag in flag_set:
        ws = wb.create_sheet(f"C{flag+1}")
        write(0, 0, "Sample Name", s="header")
        ws.column_dimensions[column_to_char(0)].width = 16
        for col, value in enumerate(dataset.classes, 1):
            write(0, col, value, s="header")
            ws.column_dimensions[column_to_char(col)].width = 10
        ws_dict[flag] = ws

    flag_index = 0
    for i, result in enumerate(results):
        row = i + 1
        if row % 2 == 0:
            style = "normal_dark"
        else:
            style = "normal_light"

        for component in result:
            flag = flags[flag_index]
            ws = ws_dict[flag]
            write(row, 0, result.sample.name, s=style)
            for col, value in enumerate(component.distribution, 1):
                write(row, col, value, s=style)
            flag_index += 1
        if progress_callback is not None:
            progress_callback(i / len(dataset) * 0.5 + 0.5)

    wb.save(filename)
    wb.close()
    if progress_callback is not None:
        progress_callback(1.0)
    logger.info(f"The SSU results have been saved to the Excel file: [{filename}].")


def save_udm(result: UDMResult, filename: str, progress_callback: Callable[[float], None] = None,
             logger: logging.Logger = None):
    if logger is None:
        logger = logging.getLogger("QGrain")
    else:
        assert isinstance(logger, logging.Logger)
    logger.debug("Start to save the UDM result.")

    wb = openpyxl.Workbook()
    prepare_styles(wb)
    logger.debug("Creating the README sheet.")
    readme_text = \
        f"""
        It contains {result.n_components+4} sheets:
        1. The first sheet is used to put the grain size distributions of corresponding samples.
        2. The second sheet is used to put the resolved parameters of all samples.
        3. The third sheet stores the statistical parameters of all components.
        4. The fourth sheet is used to put the distributions of unmixed components and the sum.
        5-{result.n_components+4}. Other sheets severally store the distributions of each component group.

        The UDM algorithm was proposed and implemented by QGrain.

        UDM Algorithm Details
            Number of Samples: {result.n_samples}
            Kernel Type: {result.kernel_type.name}
            Number of Components: {result.n_components}
            x0: {result.x0 if result.x0 is None else result.x0.tolist()}
            Number of Epochs: {len(result.loss_series("total"))}
            Spent Time: {result.time_spent} s
            
            Computing Device: {"Unknown" if result.settings is None else result.settings["device"]}
            Number of Pretrain Epochs: {"Unknown" if result.settings is None else result.settings["pretrain_epochs"]}
            Minimum Number of Epochs: {"Unknown" if result.settings is None else result.settings["min_epochs"]}
            Maximum Number of Epochs: {"Unknown" if result.settings is None else result.settings["max_epochs"]}
            Learning Rate: {"Unknown" if result.settings is None else result.settings["learning_rate"]}
            Precision: {"Unknown" if result.settings is None else result.settings["precision"]}
            Betas: {"Unknown" if result.settings is None else result.settings["betas"]}
            Constraint Level: {"Unknown" if result.settings is None else result.settings["constraint_level"]}
            Need History: {"Unknown" if result.settings is None else result.settings["need_history"]}
        """

    _write_readme_sheet(wb.active, readme_text)
    logger.debug("Creating the GSDs sheet.")
    if progress_callback is not None:
        _callback = lambda p: progress_callback(p * 0.1)
    else:
        _callback = None
    ws = wb.create_sheet("GSDs")
    _write_dataset_sheet(ws, result.dataset, progress_callback=_callback)

    def write(r, c, v, s="normal_light"):
        cell = ws.cell(r + 1, c + 1, value=v)
        cell.style = s

    logger.debug("Creating the Resolved Parameters sheet.")
    ws = wb.create_sheet("Resolved Parameters")
    write(0, 0, "Sample Name", s="header")
    ws.merge_cells(start_row=1, start_column=1, end_row=2, end_column=1)
    ws.column_dimensions[column_to_char(0)].width = 16
    headers = []
    distribution_type = DistributionType._member_map_[result.kernel_type.name]
    distribution_class = get_distribution(distribution_type)
    sub_headers = (*distribution_class.PARAMETER_NAMES, "Weight")
    for i in range(result.n_components):
        write(0, i * len(sub_headers) + 1, f"C{i+1}", s="header")
        ws.merge_cells(start_row=1, start_column=i*len(sub_headers)+2, end_row=1, end_column=(i+1)*len(sub_headers)+1)
        headers.extend(sub_headers)
    for col, value in enumerate(headers, 1):
        write(1, col, value, s="header")
        ws.column_dimensions[column_to_char(col)].width = 10
    for i, sample in enumerate(result.dataset):
        row = i + 2
        if row % 2 == 0:
            style = "normal_light"
        else:
            style = "normal_dark"
        write(row, 0, sample.name, s=style)
        for j in range(result.n_components):
            for k in range(len(sub_headers)):
                col = j * len(sub_headers) + k + 1
                value = result.parameters[-1, i, k, j]
                write(row, col, value, s=style)
        if progress_callback is not None:
            progress_callback(i / result.n_samples * 0.1 + 0.1)

    logger.debug("Creating the Statistical Moments sheet.")
    ws = wb.create_sheet("Statistical Moments")
    write(0, 0, "Sample Name", s="header")
    ws.merge_cells(start_row=1, start_column=1, end_row=2, end_column=1)
    ws.column_dimensions[column_to_char(0)].width = 16
    headers = []
    sub_headers = ["Proportion [%]",
                   "Mean [phi]",
                   "Mean [microns]",
                   "Standard Deviation [phi]",
                   "Standard Deviation [microns]",
                   "Skewness",
                   "Kurtosis"]
    for i in range(result.n_components):
        write(0, i * len(sub_headers) + 1, f"C{i+1}", s="header")
        ws.merge_cells(start_row=1, start_column=i*len(sub_headers)+2, end_row=1, end_column=(i+1)*len(sub_headers)+1)
        headers.extend(sub_headers)
    for col, value in enumerate(headers, 1):
        write(1, col, value, s="header")
        ws.column_dimensions[column_to_char(col)].width = 10

    for i, sample in enumerate(result.dataset):
        row = i + 2
        if row % 2 == 0:
            style = "normal_light"
        else:
            style = "normal_dark"
        write(row, 0, sample.name, s=style)
        for j in range(result.n_components):
            s = all_statistics(result.dataset.classes, result.dataset.classes_phi, result.components[i, j])
            write(row, j * len(sub_headers) + 1, result.proportions[i, 0, j] * 100, s=style)
            write(row, j * len(sub_headers) + 2, s["logarithmic"]["mean"], s=style)
            write(row, j * len(sub_headers) + 3, s["geometric"]["mean"], s=style)
            write(row, j * len(sub_headers) + 4, s["logarithmic"]["std"], s=style)
            write(row, j * len(sub_headers) + 5, s["geometric"]["std"], s=style)
            write(row, j * len(sub_headers) + 6, s["logarithmic"]["skewness"], s=style)
            write(row, j * len(sub_headers) + 7, s["logarithmic"]["kurtosis"], s=style)
        if progress_callback is not None:
            progress_callback(i / result.n_samples * 0.1 + 0.2)

    logger.debug("Creating the Unmixed Components sheet.")
    ws = wb.create_sheet("Unmixed Components")
    ws.merge_cells(start_row=1, start_column=1, end_row=1, end_column=2)
    write(0, 0, "Sample Name", s="header")
    ws.column_dimensions[column_to_char(0)].width = 16
    for col, value in enumerate(result.dataset.classes, 2):
        write(0, col, value, s="header")
        ws.column_dimensions[column_to_char(col)].width = 10

    predict = result.proportions @ result.components
    row = 1
    for i, sample in enumerate(result.dataset):
        if i % 2 == 0:
            style = "normal_light"
        else:
            style = "normal_dark"
        write(row, 0, sample.name, s=style)
        ws.merge_cells(start_row=row+1, start_column=1, end_row=row+result.n_components+1, end_column=1)
        for j in range(result.n_components):
            write(row, 1, f"C{j+1}", s=style)
            for col, value in enumerate(result.components[i, j]*result.proportions[i, 0, j], 2):
                write(row, col, value, s=style)
            row += 1
        write(row, 1, "Sum", s=style)
        for col, value in enumerate(predict[i, 0], 2):
            write(row, col, value, s=style)
        row += 1
        if progress_callback is not None:
            progress_callback(i / result.n_samples * 0.2 + 0.3)

    logger.debug("Creating separate sheets for all components.")
    for j in range(result.n_components):
        ws = wb.create_sheet(f"C{j+1}")
        write(0, 0, "Sample Name", s="header")
        ws.column_dimensions[column_to_char(0)].width = 16
        for col, value in enumerate(result.dataset.classes, 1):
            write(0, col, value, s="header")
            ws.column_dimensions[column_to_char(col)].width = 10
        for i, sample in enumerate(result.dataset):
            row = i + 1
            if row % 2 == 0:
                style = "normal_dark"
            else:
                style = "normal_light"
            write(row, 0, sample.name, s=style)
            for col, value in enumerate(result.components[i, j], 1):
                write(row, col, value, s=style)
            if progress_callback is not None:
                progress_callback(((i / result.n_samples) + j) / result.n_components * 0.5 + 0.5)

    wb.save(filename)
    wb.close()
    if progress_callback is not None:
        progress_callback(1.0)
    logger.info(f"The UDM result have been saved to the Excel file: [{filename}].")
