#!/usr/bin/python3

"""
Reads elemental abundance file into an array
File format:

; 10002 <--- VERSION NUMBER
; CS30336-049 <-- STAR NAME
; ApJ 222, 333 (1999) <-- SOURCE/REFERENCE
; data automatically generated by Anna Frebel <-- SOME COMMENT
; 1 <--- DATA FORMAT, USE LOG EPSILON
;	<--- NAME OF ELEMENT TO WHICH DATA IS NORMALIZED, IF ANY
; 17  <--- NUMBER OF ELEMENTS
; C		 -1.23	  0.1 <--- FORMAT="(A5,1X,F6.2,1X,F6.2)"
; N		  1.23	  0.2
; O		 -1.23	  0.3
; Fe	  1.23	  0.4
; Mg	 -1.23	  0.5
; Si	  1.23	  0.6
; Ca	 -1.23	  0.7
; Sc	  1.23	  0.8
; Ti	 -1.23	  0.9
; Cr	  1.23	 -0.1
; Mn	 -1.23	 -0.2
; Co	  1.23	 -0.3
; Ni	 -1.23	 -0.4
; Cu	  1.23	 -0.5
; Sr	 -1.23	 -0.6
; Zr	  1.23	 -0.7
; Ba	 -1.23	 -0.8
; Lo09	<--- solar reference; use "-" if not provided
;
; ; DATA FORMAT
; ; 1 log epsilon (PREFERRED DUE TO DIFFERENT SOLAR ABUNDANCES, give H!)
; ; 2 [ ]
; ; 3 [X/Y] Y= norm element, provide [Y/H] in column for Y
; ; 4 log X/Si + 6 (by number)
; ; 5 log mole fraction (REALLY PREFERRED)
; ; 6 [X/H] - assume H_sun = 7.148276E-01, norm is assumed [H]
;
; The first number is the data, the second is the error in dex.
; Use negative values to indicate upper limits; all need to be non-zero.

"""

import os
import os.path

import numpy as np

from . import DATA_DIR
from .autils.abusets import BBNAbu, SolAbu
from .autils.isotope import ion as I
from .autils.logged import Logged


# A class for each set of readings
class Star(Logged):
    """All the data read from each file"""

    def __init__(self, filename, silent=False):
        self.silent = silent
        self.setup_logger(silent=silent)

        if not os.path.isfile(filename):
            _filename = os.path.join(DATA_DIR, "stars", filename)
            if os.path.isfile(_filename):
                filename = _filename
            else:
                raise IOError(f"file {filename} not found")

        # Load BBN data for H reference
        self.BBN_data = BBNAbu(
            name=os.path.join(DATA_DIR, "ref/bbnc19.dat"), silent=silent
        )
        self._read(filename)
        # Add BBN data
        self._bbn()
        # Sort star
        self.element_abundances.sort(order="element")

    # Write list into numpy recarray
    @staticmethod
    def list_to_array(data_list):
        # Initialize numpy array of size
        data_array = np.recarray(
            (len(data_list),),
            dtype=[("element", object), ("abundance", "f8"), ("error", "f8")],
        )
        # i lines, j elements on each line
        for i, line in enumerate(data_list):
            for j, item in enumerate(line.split()):
                # Set the corresponding data type
                if j == 0:
                    data_array[i][j] = I(item)
                else:
                    data_array[i][j] = float(item)
        return data_array

    # Read file
    def _read(self, filename):
        self.setup_logger(silent=self.silent)
        with open(filename, "rt") as f:
            self.logger_file_info(f)
            # Initialize empty array for reading each line
            content = []
            # Initialize the count for number of elements
            n_elements = 9999
            # Number each line, and write into the array, stripping whitespace and carriage returns
            for i, line in enumerate(f):
                content += [line.strip()]
                if i == 0:
                    # Version number
                    self.version = int(content[0])
                    if self.version == 10000:
                        n_elements_line = 4
                    else:
                        n_elements_line = 6
                # Line 6 gives the number of elements in the file
                # Write this into n_elements
                if i == n_elements_line:
                    n_elements = int(line)
                # Read up to the last element, read one more line (solar reference) and then stop
                if i > n_elements_line + n_elements:
                    break

        # Star name
        self.name = content[1]
        n = 2
        # Additional lines for the appropriate version numbers
        if self.version >= 10001:
            # Source
            self.source = content[n]
            # Comment
            self.comment = content[n + 1]
            n += 2
        # Format
        self.data_format = int(content[n])
        n += 1
        # Normalization element
        # Data format type 6 specifies log(H/H_sun) in this line instead of a normalization element
        if self.data_format == 6:
            self.norm_element = float(content[n])
        else:
            try:
                norm = int(content[n])
            except:
                norm = content[n]
            self.norm_element = I(norm)
        n += 1
        # Number of elements
        self.n_elements = n_elements
        n += 1
        # Element abundances, the main set of data, converted to format 5
        self.element_abundances = self.list_to_array(content[n : n + n_elements])
        n += n_elements
        # Solar reference
        if self.version >= 10002:
            self.solar_ref = content[n]
            # Default solar ref
            if self.solar_ref == "":
                self.solar_ref = None
        else:
            self.solar_ref = None

        # Convert
        self.element_abundances = self._convert(
            self.element_abundances, self.data_format, self.norm_element, self.solar_ref
        )
        # Since the format has been converted, adjust the number accordingly
        self.data_format = 5

        self.close_logger(timing=" Star loaded and converted in")

    # Convert array to format 5
    def _convert(self, array, data_format, norm_element, solar_ref):
        solar_ref_file = solar_ref

        if data_format == 1:
            # Load the solar reference for Hydrogen from the Big Bang
            h_ref = self.BBN_data.Y("H")
            array.abundance[:] = np.log10((10 ** (array.abundance - 12)) * h_ref)
            return array
        elif data_format == 2:
            # Load the abundances for the sun, using the appropriate solar reference
            solar = SolAbu(solar_ref, silent=True).Y(array.element)
            array.abundance[:] = np.log10((10 ** (array.abundance)) * solar)
            return array
        elif data_format == 3:
            # Load the Big Bang hydrogen abundance
            bbn_ref_file = os.path.join(DATA_DIR, "ref/bbnc19.dat")
            h_ref = BBNAbu(bbn_ref_file, silent=True).Y("H")
            # Load the abundances for the sun
            sun = SolAbu(solar_ref_file, silent=True)
            solar = sun.Y(array.element)
            # Set the index for the norm element, since this will need to be treated seperately
            norm_index = np.where(array.element == norm_element)[0][0]
            # Calculate the linear abundance for the norm element on the star
            norm_abundance = ((10 ** array[norm_index].abundance) * h_ref) * (
                sun.Y(norm_element) / sun.Y("H")
            )
            # Calculate the abundance for the norm element on the sun
            norm_solar = sun.Y(norm_element)
            # Calculate abundances for all the other elements
            abundances = 10 ** (array.abundance) * (norm_abundance) * solar / norm_solar
            # Change the entry for the norm element that we calculated earlier
            abundances[norm_index] = norm_abundance
            # Log, and then put back into array
            array.abundance[:] = np.log10(abundances)
            return array
        elif data_format == 4:
            raise NotImplementedError()
        elif data_format == 5:
            # No changes need to be made here
            return array
        elif data_format == 6:
            # Since this is format 6, norm_element is actually log(H/H_sun)
            H = norm_element
            # Load the abundances for the sun, using the appropriate solar reference
            solar = SolAbu(solar_ref_file, silent=True).Y(array.element)
            array.abundance[:] = np.log10((10 ** (array.abundance + H)) * solar)
            return array

        else:
            raise Exception(f"Format type {data_format:d} not found")
        pass

    # A nice readable name
    def __str__(self):
        return f"{self.__class__.__name__:s}({self.name:s})"

    __repr__ = __str__

    def _bbn(self):
        """Adds BBN data to the star"""
        self.setup_logger(silent=self.silent)
        self.logger.info("Adding BBN upper limits of H and He into star data")

        BBN_abundances = np.recarray(
            (2,), dtype=[("element", object), ("abundance", "f8"), ("error", "f8")]
        )
        self.BBN_data
        for i, entry in enumerate(BBN_abundances):
            entry["element"] = I(Z=i + 1)
            entry["abundance"] = np.log10(self.BBN_data.Y(entry["element"]))
            entry["error"] = -0.1  # This is an upper limit, 0.1 is somewhat arbitrary

        # Hang onto original ones
        original_abundances = self.element_abundances
        # Make new recarray that has 2 extra places
        self.element_abundances = np.recarray(
            (2 + self.element_abundances.shape[0],),
            dtype=[("element", object), ("abundance", "f8"), ("error", "f8")],
        )

        self.element_abundances[0:2] = BBN_abundances
        self.element_abundances[2:] = original_abundances

        self.n_elements = self.element_abundances.shape[0]


def StarTable(
    filenames=[
        "Caffau.dat",
        "HE0107-5240.dat",
        "HE0557-4840.dat",
        "HE1327-2326.dat",
        "SM0313-6708.dat",
    ],
    starnames=[
        "A",
        "B",
        "C",
        "D",
        "E",
    ],
):
    stars = []
    # Read all the star files into an object, then put into list
    for filename in filenames:
        stars += [Star(filename)]

    # Use only elements that are in at least 1 star
    z = set()
    el = [] * len(stars)
    for i, star in enumerate(stars):
        el += [star.element_abundances["element"].tolist()]
        z |= {element.Z for element in el[i]}
    z = list(z)
    z = [i for i in z if i <= 30]
    z.sort()
    for elrow in el:
        elrow = [e for e in elrow if e.Z <= 30]
    maxcolumns = 7
    tablecount = round((len(z) / maxcolumns) + 0.5)
    print(r"\begin{tabular}{|l|r@{}r|r@{}r|r@{}r|r@{}r|r@{}r|r@{}r|r@{}r|}")
    for block in range(tablecount):
        print(r"\hline")
        print(r"\multicolumn{1}{|c|}{\textbf{Star}}", end="")
        for element in z[block * maxcolumns : block * maxcolumns + maxcolumns]:
            print(
                r" & \multicolumn{2}{c|}{\textbf{"
                + I(int(element)).element_symbol()
                + r"}}",
                end="",
            )
        print(r"\\")
        print(r"\hline")
        # Read those elements from the objects, then sort by charge
        for i, (star, e) in enumerate(zip(stars, el)):
            print(f"{starnames[i]:s}", end="")
            for number in z[block * maxcolumns : block * maxcolumns + maxcolumns]:
                if number in [element.Z for element in e]:
                    row = star.element_abundances[
                        [element.Z for element in e].index(number)
                    ]
                    if row.error < 0:
                        uplimstring = "<"
                    else:
                        uplimstring = " "
                    abustring = r"{:4.2f} \pm {:3.2f}".format(
                        row.abundance, abs(row.error)
                    ).rjust(15)
                    print(r" & ${:s}$ & ${:s}$".format(uplimstring, abustring), end="")
                else:
                    print(r" & $ $ & $             - $", end="")
            print(r"\\")
        print(r"\hline")
        if block == tablecount:
            break
        else:
            print(r"\noalign{\smallskip}")
    print(r"\end{tabular}")


if __name__ == "__main__":
    import sys

    args = sys.argv[1:]
    for i, a in enumerate(args):
        try:
            args[i] = eval(a)
        except:
            pass
    print(Star(*args))
