import os
import copy
import types
import xlwt
import numpy
import base64
import traceback
from json import JSONEncoder
import json
from ..ds_ccd import DiffImg
from ..ds_jcpds import JCPDSplt, Session
from ..ds_ramspec import PatternPeakPo, get_data_section
from ..ds_section import Section
from .diff_state import DiffState
from ..utils import samefilename, make_filename, change_file_path, \
    extract_extension


class PeakPoModel(object):
    """
    session is only for reading/writing/referencing.
    components of the models are not part of session.
    session is a reference object
    """

    def __init__(self):
        self.base_ptn = None
        self.waterfall_ptn = []
        self.jcpds_lst = []
        self.diff_img = None
        self.poni = None
        self.session = None
        self.jcpds_path = ''
        self.chi_path = ''
        self.current_section = None
        self.section_lst = []
        self.saved_pressure = 10.
        self.saved_temperature = 300.

    def exist_in_waterfall(self, filename):
        if not self.waterfall_exist():
            return False
        for f in self.waterfall_ptn:
            if filename == f.fname:
                return True
        return False

    def get_saved_pressure(self):
        return self.saved_pressure

    def get_saved_temperature(self):
        return self.saved_temperature

    def save_pressure(self, pressure):
        self.saved_pressure = pressure

    def save_temperature(self, temperature):
        self.saved_temperature = temperature

    def set_this_section_current(self, index):
        self.current_section = None
        self.current_section = copy.deepcopy(self.section_lst[index])

    def clear_section_list(self):
        self.section_lst[:] = []

    def get_number_of_section(self):
        return self.section_lst.__len__()

    def set_current_section(self, roi):
        x_section_bg, y_section_bg = get_data_section(
            self.base_ptn.x_bg, self.base_ptn.y_bg, roi)
        __, y_section_bgsub = get_data_section(
            self.base_ptn.x_bgsub, self.base_ptn.y_bgsub, roi)
        self.current_section.set(x_section_bg, y_section_bgsub, y_section_bg)

    def current_section_exists_in_list(self):
        for section in self.section_lst:
            if self.current_section.get_timestamp() == section.get_timestamp():
                return True
        return False

    def current_section_saved(self):
        if self.get_number_of_section() == 0:
            return False
        if self.current_section.timestamp is None:
            return False
        if self.current_section_exists_in_list():
            return True
        else:
            return False

    def initialize_current_section(self):
        if self.current_section_exist():
            self.current_section = None
        self.current_section = Section()

    def save_current_section(self):
        new_section = copy.deepcopy(self.current_section)
        self.section_lst.append(new_section)
        self.current_section = None

    def current_section_exist(self):
        if self.current_section is None:
            return False
        if self.current_section.x is None:
            return False
        else:
            return True

    def set_from(self, model_r, new_chi_path=None, jlistonly=False):
        self.jcpds_lst = model_r.jcpds_lst
        if jlistonly:
            return
        if new_chi_path is None:
            pass
        else:
            new_base_ptn_fname = change_file_path(model_r.base_ptn.fname,
                                                  new_chi_path)
            model_r.base_ptn.fname = new_base_ptn_fname
            if model_r.waterfall_ptn != []:
                new_waterfall_ptn = []
                for ptn in model_r.waterfall_ptn:
                    new_fname = change_file_path(ptn.fname, new_chi_path)
                    if os.path.exists(new_fname):
                        ptn.fname = new_fname
                        new_waterfall_ptn.append(ptn)
                if new_waterfall_ptn != []:
                    model_r.waterfall_ptn = new_waterfall_ptn
            if model_r.diff_img is not None:
                new_img_fname = change_file_path(
                    model_r.diff_img.img_filename,
                    new_chi_path)
                model_r.diff_img.img_filename = new_img_fname
            if model_r.poni is not None:
                new_poni_fname = change_file_path(model_r.poni,
                                                  new_chi_path)
                model_r.poni = new_poni_fname
            model_r.chi_path = new_chi_path
        self.base_ptn = model_r.base_ptn
        self.waterfall_ptn = model_r.waterfall_ptn
        self.diff_img = model_r.diff_img
        self.section_lst = model_r.section_lst
        self.saved_pressure = model_r.get_saved_pressure()
        self.saved_temperature = model_r.get_saved_temperature()
        self.poni = model_r.poni
        self.session = model_r.session
        self.jcpds_path = model_r.jcpds_path
        self.chi_path = model_r.chi_path
        # added 2023.03.26
        self.current_section = None

    def import_section_list(self, model_r):
        new_section_lst = copy.deepcopy(model_r.section_lst)
        if new_section_lst == []:
            return
        for section in new_section_lst:
            section.invalidate_fit_result()
            roi = section.get_xrange()
            x, y_bgsub, y_bg = self.get_single_section(roi)
            section.set(x, y_bgsub, y_bg)
        existing_section_lst = copy.deepcopy(self.section_lst)
        self.section_lst[:] = []
        self.section_lst = existing_section_lst + new_section_lst

    def get_single_section(self, roi):
        x_section_bg, y_section_bg = get_data_section(
            self.base_ptn.x_bg, self.base_ptn.y_bg, roi)
        __, y_section_bgsub = get_data_section(
            self.base_ptn.x_bgsub, self.base_ptn.y_bgsub, roi)
        return x_section_bg, y_section_bgsub, y_section_bg

    def reset_base_ptn(self):
        self.base_ptn = PatternPeakPo()

    def reset_waterfall_ptn(self):
        self.waterfall_ptn[:] = []

    def reset_jcpds_lst(self):
        self.jcpds_lst[:] = []

    def reset_diff_img(self):
        self.diff_img = DiffImg()

    def reset_poni(self):
        self.poni = None

    def base_ptn_exist(self):
        if self.base_ptn is None:
            return False
        else:
            if self.base_ptn.fname is None:
                return False
            else:
                return True

    def waterfall_exist(self):
        if self.waterfall_ptn == []:
            return False
        else:
            return True

    def jcpds_exist(self):
        if self.jcpds_lst == []:
            return False
        else:
            return True

    def diff_img_exist(self):
        if self.diff_img is None:
            return False
        else:
            return True

    def poni_exist(self):
        if self.poni is None:
            return False
        else:
            return True

    def make_filename(self, extension, original=False):
        """
        :param extension: extension without a dot
        """
        return make_filename(self.base_ptn.fname, extension, original=original)

    def get_associated_image_candidates(self):
        """
        Candidate raw-image filenames expected for current CHI.
        Includes both legacy (original=True) and exact-stem naming.
        """
        if self.base_ptn_exist() and \
                os.path.splitext(str(self.base_ptn.fname))[1].lower() == ".spe":
            return [str(self.base_ptn.fname)]
        exts = ("tif", "tiff")
        candidates = []
        for ext in exts:
            for original in (True, False):
                filen = self.make_filename(ext, original=original)
                if filen not in candidates:
                    candidates.append(filen)
        return candidates

    def same_filename_as_base_ptn(self, filename):
        return samefilename(self.base_ptn.fname, filename)

    def set_base_ptn(self, new_base_ptn_filen, wavelength):
        """
        :param new_base_ptn: PatternPeakPo object
        """
        self.reset_base_ptn()
        self.base_ptn.read_file(new_base_ptn_filen)
        self.set_chi_path(os.path.split(new_base_ptn_filen)[0])
        self.set_base_ptn_wavelength(wavelength)
        self.base_ptn.display = True

    def get_base_ptn(self):
        return self.base_ptn

    def append_a_waterfall_ptn(self, filename, wavelength,
                               bg_roi, bg_params, temp_dir=None):
        pattern = PatternPeakPo()
        pattern.read_file(filename)
        pattern.wavelength = wavelength
        if hasattr(pattern, "apply_excitation_wavelength"):
            pattern.apply_excitation_wavelength(wavelength)
        pattern.display = False
        if temp_dir is None:
            pattern.get_chbg(bg_roi, params=bg_params, yshift=0)
        else:
            success = pattern.read_bg_from_tempfile(temp_dir=temp_dir)
            if not success:
                pattern.get_chbg(bg_roi, params=bg_params, yshift=0)
        self.waterfall_ptn.append(pattern)

    def replace_a_waterfall(self, new_pattern, index_to_replace):
        self.waterfall_ptn[index_to_replace] = new_pattern

    def set_waterfall_ptn(
            self, filenames, wavelength, display, bg_roi, bg_params,
            temp_dir=None):
        new_waterfall_ptn = []
        for f, wl, dp in zip(filenames, wavelength, display):
            pattern = PatternPeakPo()
            pattern.read_file(f)
            pattern.wavelength = wl
            if hasattr(pattern, "apply_excitation_wavelength"):
                pattern.apply_excitation_wavelength(wl)
            pattern.display = dp
            if temp_dir is None:
                pattern.get_chbg(bg_roi, params=bg_params, yshift=0)
            else:
                success = pattern.read_bg_from_tempfile(temp_dir=temp_dir)
                if not success:
                    pattern.get_chbg(bg_roi, params=bg_params, yshift=0)
            new_waterfall_ptn.append(pattern)
        self.waterfall_ptn = new_waterfall_ptn

    def append_a_jcpds(self, filen, color, cif_k0=200.0, cif_k0p=4.0,
                       cif_alpha=1e-5):
        try:
            phase = JCPDSplt()
            ext = extract_extension(filen).lower()
            if ext == 'cif':
                cif_name = os.path.splitext(os.path.basename(filen))[0]
                success = phase.set_from_cif(
                    filen, cif_k0, cif_k0p, file=filen, name=cif_name,
                    comments='Created from ' + filen,
                    thermal_expansion=cif_alpha)  # phase.file = f
                if not success:
                    return False
            else:
                phase.read_file(filen)
            phase.color = color
        except Exception as e:
            print("append_a_jcpds failed for:", filen)
            print("error:", repr(e))
            traceback.print_exc()
            return False
        self.jcpds_lst.append(phase)
        return True

    def write_as_ppss(self, fname, pressure, temperature):
        raise NotImplementedError("PPSS export was removed. Use the JSON session folder.")

    def read_ppss(self, fname):
        raise NotImplementedError("PPSS import was removed. Use the JSON session folder.")

    def set_jcpds_from_ppss(self):
        raise NotImplementedError("PPSS import was removed. Use the JSON session folder.")

    def set_chi_path(self, chi_path):
        self.chi_path = chi_path

    def set_jcpds_path(self, jcpds_path):
        self.jcpds_path = jcpds_path

    def get_base_ptn_wavelength(self):
        return self.base_ptn.wavelength

    def set_base_ptn_wavelength(self, wavelength):
        self.base_ptn.wavelength = wavelength
        if hasattr(self.base_ptn, "apply_excitation_wavelength"):
            self.base_ptn.apply_excitation_wavelength(wavelength)
        if self.diff_img is not None and hasattr(self.diff_img, "apply_excitation_wavelength"):
            self.diff_img.apply_excitation_wavelength(wavelength)

    def get_base_ptn_filename(self):
        return self.base_ptn.fname

    def set_base_ptn_color(self, color):
        self.base_ptn.color = color

    def associated_image_exists(self):
        if self.base_ptn_exist() and \
                os.path.splitext(str(self.base_ptn.fname))[1].lower() == ".spe":
            return True
        for filen in self.get_associated_image_candidates():
            if os.path.exists(filen):
                return True
        return False

    def load_associated_img(self):
        self.reset_diff_img()
        if self.base_ptn_exist() and \
                os.path.splitext(str(self.base_ptn.fname))[1].lower() == ".spe":
            self.diff_img.load(self.base_ptn.fname)
            if hasattr(self.diff_img, "apply_excitation_wavelength"):
                self.diff_img.apply_excitation_wavelength(self.base_ptn.wavelength)
            return True
        for filen_toload in self.get_associated_image_candidates():
            if os.path.exists(filen_toload):
                self.diff_img.load(filen_toload)
                if hasattr(self.diff_img, "apply_excitation_wavelength"):
                    self.diff_img.apply_excitation_wavelength(self.base_ptn.wavelength)
                return True
        return False

    def section_list_exist(self):
        if self.section_lst == []:
            return False
        else:
            return True

    def save_peak_fit_results_to_xls(self, xls_filen):
        """
        returns boolean for success
        """
        if not self.section_list_exist():
            return False
        if str(xls_filen) == '':
            return
        num_sec = 0
        workbook = xlwt.Workbook()
        sheet_num = 0
        for section in self.section_lst:
            x_range = section.get_xrange()
            xmin = x_range[0]
            xmax = x_range[1]
            sheet_name = "{0:d}_at_{1:.2f}-{2:.2f}".format(
                sheet_num, xmin, xmax)
            sheet_num += 1
            sheet = workbook.add_sheet(sheet_name)
            sheet.write(0, 0, section.timestamp)
            sheet.write(1, 0, 'Pressure (GPa)')
            sheet.write(2, 0, 'Temperature (K)')
            sheet.write(1, 1, self.get_saved_pressure())
            sheet.write(2, 1, self.get_saved_temperature())
            sheet.write(3, 0, 'Section x range')
            sheet.write(3, 1, xmin)
            sheet.write(3, 2, xmax)
            sheet.write(4, 0, 'Chisqr')
            sheet.write(4, 1, section.fit_result.chisqr)
            sheet.write(5, 0, 'Reduced Chisqr')
            sheet.write(5, 1, section.fit_result.redchi)
            sheet.write(6, 0, 'Akaike info crit')
            sheet.write(6, 1, section.fit_result.aic)
            sheet.write(7, 0, 'Bayesian info crit')
            sheet.write(7, 1, section.fit_result.bic)
            # write peak params and errors first
            lineno = 8
            sheet.write(lineno, 1, 'Phase')
            sheet.write(lineno, 2, 'h')
            sheet.write(lineno, 3, 'k')
            sheet.write(lineno, 4, 'l')
            sheet.write(lineno, 5, 'Area value')
            sheet.write(lineno, 6, 'Area stderr')
            sheet.write(lineno, 7, 'Area vary')
            sheet.write(lineno, 8, 'Pos value')
            sheet.write(lineno, 9, 'Pos stderr')
            sheet.write(lineno, 10, 'Pos vary')
            sheet.write(lineno, 11, 'FWHM value')
            sheet.write(lineno, 12, 'FWHM stderr')
            sheet.write(lineno, 13, 'FWHM vary')
            sheet.write(lineno, 14, 'nL value')
            sheet.write(lineno, 15, 'nL stderr')
            sheet.write(lineno, 16, 'nL vary')
            lineno += 1
            n_peak = section.get_number_of_peaks_in_queue()
            for i in range(n_peak):
                prefix = "p{0:d}_".format(i)
                sheet.write(lineno, 0, prefix)
                sheet.write(lineno, 1, section.peakinfo[prefix + 'phasename'])
                sheet.write(lineno, 2, section.peakinfo[prefix + 'h'])
                sheet.write(lineno, 3, section.peakinfo[prefix + 'k'])
                sheet.write(lineno, 4, section.peakinfo[prefix + 'l'])
                sheet.write(lineno, 5, section.fit_result.
                            params[prefix + 'amplitude'].value)
                sheet.write(lineno, 6, section.fit_result.
                            params[prefix + 'amplitude'].stderr)
                sheet.write(lineno, 7, section.fit_result.
                            params[prefix + 'amplitude'].vary)
                sheet.write(lineno, 8, section.fit_result.
                            params[prefix + 'center'].value)
                sheet.write(lineno, 9, section.fit_result.
                            params[prefix + 'center'].stderr)
                sheet.write(lineno, 10, section.fit_result.
                            params[prefix + 'center'].vary)
                sheet.write(lineno, 11, section.fit_result.
                            params[prefix + 'sigma'].value * 2.)
                if section.fit_result.params[prefix + 'sigma'].stderr is None:
                    sheet.write(lineno, 12, 'None')
                else:
                    sheet.write(lineno, 12, section.fit_result.
                                params[prefix + 'sigma'].stderr * 2.)
                sheet.write(lineno, 13, section.fit_result.
                            params[prefix + 'sigma'].vary)
                sheet.write(lineno, 14, section.fit_result.
                            params[prefix + 'fraction'].value)
                sheet.write(lineno, 15, section.fit_result.
                            params[prefix + 'fraction'].stderr)
                sheet.write(lineno, 16, section.fit_result.
                            params[prefix + 'fraction'].vary)
                lineno += 1
            lineno += 1
            sheet.write(lineno, 0, 'Baseline factors')
            lineno += 1
            n_order = section.get_order_of_baseline_in_queue()
            sheet.write(lineno, 1, 'value')
            sheet.write(lineno, 2, 'stderr')
            sheet.write(lineno, 3, 'vary')
            lineno += 1
            for i in range(n_order + 1):
                prefix = "b_c{0:d}".format(i)
                sheet.write(lineno, 0, prefix)
                sheet.write(lineno, 1, section.fit_result.params[prefix].value)
                sheet.write(lineno, 2, section.fit_result.params[prefix].stderr)
                sheet.write(lineno, 3, section.fit_result.params[prefix].vary)
                lineno += 1
            lineno += 2
            sheet.write(lineno, 0, 'x_data')
            sheet.write(lineno, 1, 'y_data')
            sheet.write(lineno, 2, 'y_bgsub')
            sheet.write(lineno, 3, 'y_bg')
            sheet.write(lineno, 4, 'y_fit_profile')
            y_total_profile = section.get_fit_profile(bgsub=False)
            y_single_profiles = section.get_individual_profiles(bgsub=False)
            k = 0
            for key, value in y_single_profiles.items():
                sheet.write(lineno, 5 + k, key + 'profile')
                k += 1
            lineno += 1
            for i in range(section.x.__len__()):
                sheet.write(lineno, 0, section.x[i])
                sheet.write(lineno, 1, section.y_bg[i] + section.y_bgsub[i])
                sheet.write(lineno, 2, section.y_bgsub[i])
                sheet.write(lineno, 3, section.y_bg[i])
                sheet.write(lineno, 4, y_total_profile[i])
                j = 0
                for key, value in y_single_profiles.items():
                    sheet.write(lineno, 5 + j, value[i])
                    j += 1
                lineno += 1
        workbook.save(xls_filen)


class PeakPoModel707(PeakPoModel):
    def __init__(self):
        """
        azisections: azimuthal section object. inherit base_ptn object
            azimutal angles, and all pattern variables
            wonder if I need to revise the fit sections to link to Azimuthal
            sections
        """
        self.azisections = []

class PeakPoModel8(PeakPoModel):
    def __init__(self):
        self.base_ptn = None
        self.waterfall_ptn = []
        self.jcpds_lst = []
        self.diff_img = None
        self.poni = None
        self.session = None
        self.jcpds_path = ''
        self.chi_path = ''
        self.current_section = None
        self.section_lst = []
        self.saved_pressure = 10.
        self.saved_temperature = 300.
        self.dum = None
        self.diff_state = DiffState()

    def from_model7(self, model7):
        self.base_ptn = model7.base_ptn
        self.waterfall_ptn = model7.waterfall_ptn
        self.jcpds_lst = model7.jcpds_lst
        self.diff_img = model7.diff_img
        self.poni = model7.poni
        self.session = model7.session
        self.jcpds_path = model7.jcpds_path
        self.chi_path = model7.chi_path
        self.current_section = model7.current_section
        self.section_lst = model7.section_lst
        self.saved_pressure = model7.saved_pressure
        self.saved_temperature = model7.saved_temperature
        self.diff_state = DiffState()

    def to_model7(self):
        model7 = PeakPoModel()
        model7.base_ptn = self.base_ptn
        model7.waterfall_ptn = self.waterfall_ptn
        model7.jcpds_lst = self.jcpds_lst
        model7.diff_img = self.diff_img
        model7.poni = self.poni
        model7.session = self.session
        model7.jcpds_path = self.jcpds_path
        model7.chi_path = self.chi_path
        model7.current_section = self.current_section
        model7.section_lst = self.section_lst
        model7.saved_pressure = self.saved_pressure
        model7.saved_temperature = self.saved_temperature
        return model7

    def save_to_txtdata(self, temp_dir):
        # base_ptn
        write_JSON(self.base_ptn, 'base_ptn', temp_dir)
        # waterfall_ptn
        write_JSON(self.waterfall_ptn, 'waterfall_ptn', temp_dir)
        # jcpds_list
        write_JSON(self.jcpds_lst, 'jcpds_lst', temp_dir)
        # diff_img
        if self.diff_img != None:
            dict = diffimg_to_dict(self.diff_img)
            json_str = json.dumps(dict, indent=4, cls=PeakPoEncoder)
            json_filen = os.path.join(temp_dir, 'diff_img.json')
            with open(json_filen, 'w') as outfile:
                outfile.write(json_str)
        # poni
        write_JSON(self.poni, 'poni', temp_dir)
        # session
        write_JSON(self.session, 'session', temp_dir)
        # jcpds_path
        write_JSON(self.jcpds_path, 'jcpds_path', temp_dir)
        # chi_path
        write_JSON(self.chi_path, 'chi_path', temp_dir)
        # current_section
        write_JSON(self.current_section, 'current_section', temp_dir)
        # section_lst
        # need to solve below.  temporary blocked.
        # the reason is that section_lst is collection of Section object 
        # which is not serialized yet.
        if self.section_lst is None:
            print('self section_lst none')
        else:
            print(self.section_lst)
        # write_JSON(self.section_lst, 'section_lst', temp_dir)
        # saved_pressure
        write_JSON(self.saved_pressure, 'saved_pressure', temp_dir)
        # saved_temperature
        write_JSON(self.saved_temperature, 'saved_temperature', temp_dir)

    def load_from_txtdata(self):
        # https://pynative.com/python-convert-json-data-into-custom-python-object/
        print('I am here')

def write_JSON(obj, filename, dir):
    json_str = json.dumps(obj, indent=4, cls=PeakPoEncoder)
    json_filen = os.path.join(dir, filename+'.json')
    with open(json_filen, 'w') as outfile:
        outfile.write(json_str)

# https://stackoverflow.com/questions/27909658/json-encoder-and-decoder-for-complex-numpy-arrays
class PeakPoEncoder(JSONEncoder):
    def default(self, obj):
        if isinstance(obj, numpy.ndarray):
            #data_b64 = base64.b64encode(obj.data)
            #return dict(__ndarray__=data_b64,
            #            dtype=str(obj.dtype),
            #            shape=obj.shape)
            return obj.tolist()
        elif isinstance(obj, types.MappingProxyType):
            return dict(obj)
        elif hasattr(obj, "__dict__"):
            return obj.__dict__
        else:
            return str(obj)
        
class PeakPoEncoder1(JSONEncoder):
    def default(self, obj):
        if isinstance(obj, numpy.ndarray):
            #data_b64 = base64.b64encode(obj.data)
            #return dict(__ndarray__=data_b64,
            #            dtype=str(obj.dtype),
            #            shape=obj.shape)
            return None
        else:
            return obj.__dict__
        
def check_NONE(item):
    if item is None:
        return None
    else:
        return item.tolist()
        
def diffimg_to_dict(diffimg):
    dict = {'img_filename': diffimg.img_filename, 'mask': check_NONE(diffimg.mask)}
    return dict
