import logging
from collections import OrderedDict

import numpy as np

from aspire.storage import StarFile

logger = logging.getLogger(__name__)


# The metadata_fields dictionary below specifies default data types
# of certain key fields used in the codebase,
# which are originally read from Relion STAR files.
relion_metadata_fields = {
    "_rlnVoltage": float,
    "_rlnDefocusU": float,
    "_rlnDefocusV": float,
    "_rlnDefocusAngle": float,
    "_rlnSphericalAberration": float,
    "_rlnDetectorPixelSize": float,
    "_rlnCtfFigureOfMerit": float,
    "_rlnMagnification": float,
    "_rlnAmplitudeContrast": float,
    "_rlnImageName": str,
    "_rlnOriginalName": str,
    "_rlnCtfImage": str,
    "_rlnCoordinateX": float,
    "_rlnCoordinateY": float,
    "_rlnCoordinateZ": float,
    "_rlnNormCorrection": float,
    "_rlnMicrographName": str,
    "_rlnGroupName": str,
    "_rlnGroupNumber": str,
    "_rlnOriginX": float,
    "_rlnOriginY": float,
    "_rlnAngleRot": float,
    "_rlnAngleTilt": float,
    "_rlnAnglePsi": float,
    "_rlnClassNumber": int,
    "_rlnLogLikeliContribution": float,
    "_rlnRandomSubset": int,
    "_rlnParticleName": str,
    "_rlnOriginalParticleName": str,
    "_rlnNrOfSignificantSamples": float,
    "_rlnNrOfFrames": int,
    "_rlnMaxValueProbDistribution": float,
    "_rlnOpticsGroup": int,
    "_rlnOpticsGroupName": str,
}


def dict_to_relion_types(d):
    """
    Convert STAR file strings to data type for each field in a dict loaded via
    `aspire.storage.StarFile`. Columns without a specified data type are read as
    `dtype=str`.

    :param d: A `StarFile` block represented as a dictionary.
    :return: A dict with types converted where possible.
    """
    column_types = {name: relion_metadata_fields.get(name, str) for name in d}
    # look at values to detect if we're dealing with iterables
    multiple_rows = isinstance(d, dict) and all(
        not isinstance(v, str) and hasattr(v, "__iter__") for v in d.values()
    )
    if multiple_rows:
        retval = dict(zip(d, [np.array(list(map(column_types[k], d[k]))) for k in d]))
    else:
        retval = {k: column_types[k](d[k]) for k in d}
    return retval


class RelionStarFile(StarFile):
    """
    A star file generated by RELION representing particles, micrographs, or movies.
    """

    def __init__(self, filepath):
        super().__init__(filepath, blocks=None)

        # convert dtypes in star file blocks where possible
        self._convert_dtypes()

        # validate star file and detect Relion version (3.0 or >=3.1)
        self._validate_and_detect_version()

    def _validate_and_detect_version(self):
        """
        Based on the number of StarFile blocks, block names, and column names,
        determine whether this star file can be interpreted as coming from RELION.
        If so, determine what version of RELION and what kind of data is represented.
        """

        self.data_block_name = ""
        self.relion_version = ""

        # validate 3.0 STAR file
        if len(self.blocks) == 1:
            data_block = self.get_block_by_index(0)
            columns = list(data_block.keys())
            if not any(
                col in columns
                for col in ["_rlnImageName", "_rlnMicrographName", "_rlnMovieName"]
            ):
                raise ValueError(
                    f"{self.filepath} does not contain Relion data columns."
                )

            self.relion_version = "3.0"
            self.data_block = data_block

        # validate 3.1 STAR file
        if len(self.blocks) == 2:
            # must have an optics block
            if "optics" not in self.blocks.keys():
                raise ValueError(f"{self.filepath} does not contain an optics block.")

            # find type of data
            for name in self.blocks.keys():
                if name in ["particles", "micrographs", "movies"]:
                    self.data_block_name = name
                    break
            if not self.data_block_name:
                raise ValueError(
                    f"{self.filepath} does not contain a block identifying particle, ",
                    "micrograph, or movie data.",
                )

            data_block = self[self.data_block_name]
            # lastly, data block must contain a column identifying the type of data as well
            columns = list(data_block.keys())
            if not any(
                col in columns
                for col in ["_rlnImageName", "_rlnMicrographName", "_rlnMovieName"]
            ):
                raise ValueError(
                    f"{self.filepath} data block does not contain Relion data columns."
                )

            self.relion_version = "3.1"
            self.data_block = data_block
            self.optics_block = self["optics"]

        if not self.relion_version:
            raise ValueError(
                f"Cannot interpret {self.filepath} as a valid RELION STAR file."
            )

    def _convert_dtypes(self):
        """
        For data blocks, convert known Relion columns to sensible types.
        """
        _blocks = OrderedDict()
        for block_name, block in self.blocks.items():
            _blocks[block_name] = dict_to_relion_types(block)
        self.blocks = _blocks

    def get_merged_data_block(self):
        """
        Return the dictionary containing particle/micrograph/movie information for this STAR file.

        :return: A dictionary containing the data from this RELION STAR file.
        """

        if self.relion_version == "3.0":
            # 3.0 star files have the data all in one block, which we return
            return self.data_block

        if self.relion_version == "3.1":
            # merge the parameters in the optics block as new columns in the data block
            # based on the corresponding optics group number (returns a new dataframe)
            data_block = self.data_block.copy()
            # Create a dict from optics_group to individual optics properties
            optics = self.optics_block
            optics_props_by_group_id = dict(
                (k, {_k: optics[_k][i] for _k in optics})
                for i, k in enumerate(optics["_rlnOpticsGroup"])
            )
            # get a NumPy array of optics indices for each row of data
            optics_indices = self.data_block["_rlnOpticsGroup"].astype(int)
            for k in optics:
                data_block[k] = np.array(
                    [
                        optics_props_by_group_id[optic_index][k]
                        for optic_index in optics_indices
                    ]
                )

            return data_block
