import numpy as np
import matplotlib.pyplot as plt
import collections

from py4xs.data2d import Data2d
from py4xs.mask import Mask
from py4xs.utils import common_name,smooth

import os
import sys
import time
import copy
import itertools as it
import multiprocessing as mp

# this is the ratio between protein average denstiy and water density
# it is assumed to be a constant but in reality depends on the specific portein
# see Fischer et.al.  Protein Sci. 2004 October; 13(10): 2825-2828
PROTEIN_WATER_DENSITY_RATIO = 1.35

# Each Data1d corresponds to one single scattering pattern
# The intensity is normalized based on
#     (1) beam intensity through the beam stop, as in prvious version
#  or (2) WAXS intensity (water scattering)
#  or (3) an externally dtermined value
# The intensity can be further normalized a reference trans value, so
#   that different sets can be compared.
# Data1d sets (must share the same qgrid) can be merged: e.g. SAXS and WAXS
# Background subtraction and flat field correction are also supported

TRANS_EXTERNAL = 0
TRANS_FROM_BEAM_CENTER = 1
TRANS_FROM_WAXS = 2

from enum import Enum
# removed from_beam_center since it is really external
# also it is difficult to keep track when merging multiple detectors
class trans_mode(Enum):
    external = 0
    from_waxs = 2

# trans_mode=TRANS_FROM_BEAM_CENTER
# this works if there is a semi-transparent beam stop
BEAM_SIZE_hW = 5
BEAM_SIZE_hH = 4

# this is the global setting
TRANS_MODE = trans_mode.from_waxs
# this is the minimum intensity to be used for trans calculations
WAXS_THRESH = 10

# this is the scaling factor for indivudual curves that belong to the same sample
# they are offset for clarity in the plots
VOFFSET = 1.5

font_size_list = ['xx-small', 'x-small', 'small', 'medium', 'large', 'x-large', 'xx-large']
def get_font_size(size_index):
    """ "medium" has size_index of 0
        the size_index is negative for smaller fonts and possitive for larger ones
    """
    if size_index in font_size_list:
        i = font_size_list.index(size_index)
    else:
        i = int(size_index)+3
        if i<0:
            i = 0
        elif i>=len(font_size_list):
            i = len(font_size_list)-1
    return i-3,font_size_list[i]

class Data1d:
    def __init__(self, transMode=None):
        self.comments = ""
        self.label = "data"
        self.overlaps = []
        self.raw_data = {}
        self.timestamp = None
        self.trans_w = -1
        self.trans_e = -1
        self.trans = -1
        self.transMode = transMode

    def load_from_2D(self, image, exp_para, qgrid, pre_process=None, flat_cor=None,
                     mask=None, save_ave=False, debug=False, label=None, dtype=None):
        """
        image: a filename, or a Data2d instance, or a numpy array
        qgrid: for the 1D data
        exp_para: ExpPara
        mask: no longer used, extract from exp_para
        """
        self.qgrid = qgrid
        mask = exp_para.mask

        if debug==True:
            print("loading data from 2D image: ", label)

        if isinstance(image, Data2d):
            d2 = image
        else:
            d2 = Data2d(image, exp=exp_para, dtype=dtype)
            self.timestamp = d2.timestamp
            self.label = d2.label
            self.timestamp = d2.timestamp

        if label is not None:
            self.label = label

        # place holder for pre-processing, to deal with things like
        # dark current, flat field, and dezinger corrections on the 2D data
        if pre_process is not None:
            pre_process(d2.data)

        cor_factor = exp_para.FSA*exp_para.FPol
        if flat_cor is not None:
            cor_factor *= flat_cor    # for rescuing data with incorrect flat field at the time of collection
        self.data,self.err = d2.conv_Iq(qgrid, mask, cor_factor = cor_factor)

        if isinstance(image, np.ndarray):
            del d2      # d2 is only used temporarily

        if save_ave and isinstance(image, str):
            self.save(image + ".ave", debug=debug)

    def set_trans(self, transMode=None, trans=-1, ref_trans=-1,
                  calc_water_peak=False, q_start=1.85, q_end=2.15, debug=False):
        """
        normalize intensity, from trans to ref_trans
        trans can be either from the beam center/beam stop or water scattering
        sometimes the data may require both measures of transmitted intensity

        this operation should be performed after SAXS/WAXS merge, because
          1. SAXS and WAXS should have the same trans
          2. if trans_mode is TRNAS_FROM_WAXS, the trans value needs to be calculated from WAXS data

        """
        if transMode is None and self.transMode is not None:
            transMode = self.transMode
        else:
            assert(isinstance(transMode, trans_mode))
            self.transMode = transMode
        if self.transMode==trans_mode.from_waxs or calc_water_peak:
            # get trans for the near the maximum in the WAXS data
            # for solution scattering, hopefully this reflect the intensity of water scattering
            idx = (self.qgrid > q_start) & (self.qgrid < q_end)  # & (self.data>0.5*np.max(self.data))
            if len(self.qgrid[idx]) < 5:
                print("not enough data points under the water peak, consider using a different trans_mode.")
                #raise Exception()

            # trying to narrow down the peak range turns out to be a bad idea
            # the width then could vary between datasets, creating artificial fluctuation in trans
            #idx1 = idx & (self.data >= 0.95*np.max(self.data[idx]))

            if (self.data[idx]<WAXS_THRESH).all() and debug!='quiet':
                print("the data points for trans calculation are below WAXS_THRESH: ",
                      np.max(self.data[idx]), WAXS_THRESH)
            self.trans_w = np.sum(self.data[idx])
            if self.transMode==trans_mode.from_waxs:
                self.trans = self.trans_w
            qavg = np.average(self.qgrid[idx])
            if self.trans_w<1.0:
                print(f'caluclated trans is {self.trans_w}, setting it artifically to WAXS_THRESH.')
                self.trans_w = WAXS_THRESH
            if debug==True:
                print("using data near the high q end (q~%f)" % qavg, end=' ')
            self.comments += "# transmitted beam intensity from WAXS (q~%.2f)" % qavg
        if self.transMode==trans_mode.external and trans>=0:
            # if the trans value is specified, by definition transMode should be external
            # if trans=-1, it is not meant to be the value to be set
            #if trans<0:
            #    print(f"Warning: {trans} is not a valid value for transmitted intensity.")
            #    trans = 0
            self.comments += f"# transmitted beam intensity given externally: {trans}"
            self.trans_e = trans
            self.trans = trans
            #if self.transMode==trans_mode.external:
            #    self.trans = trans

        self.comments += ": %f \n" % self.trans
        if debug==True:
            print("trans for %s set to %f" % (self.label, self.trans))

        if ref_trans > 0:
            if self.trans<0:
                print(f"cannot normalize intensity since data1d does not have a valid trans value: {self.trans}")
            self.comments += "# scattering intensity normalized to ref_trans = %f \n" % ref_trans
            self.data *= ref_trans/self.trans
            self.err *= ref_trans/self.trans
            for ov in self.overlaps:
                ov['raw_data1'] *= ref_trans/self.trans
                ov['raw_data2'] *= ref_trans/self.trans
            self.trans_w *= ref_trans/self.trans
            self.trans_e *= ref_trans/self.trans
            self.trans = ref_trans
            if debug==True:
                print("normalized to %f" % ref_trans)

    def avg(self, dsets, weighted=True, qmax_for_weight=0.3,
            plot_data=False, ax=None, debug=False, fontsize='large'):
        """
        dsets is a collection of Data1d
        weighted:
            if False
                the dsets are simply averaged together
                errorbar is increased if there are discrepencies between individual values??
            otherwise
                weight should contain a list of weight factors, corresponding to each dset
                each dset is first scaled by the
                a weighted average (smaller errorbar has higher weight) is then performed
        ax is the Axes to plot the data in
        TODO: should calculate something like the cross-correlation between sets
        to evaluate the consistency between them
        """
        if debug!='quiet':
            print("averaging data with %s: \n" % self.label, end=' ')
        i_fs = get_font_size(fontsize)[0]

        n = 1
        if plot_data:
            if ax is None:
                plt.figure()
                plt.subplots_adjust(bottom=0.15)
                ax = plt.gca()
            ax.set_xlabel("$q (\AA^{-1})$", fontsize=get_font_size(i_fs)[1])
            ax.set_ylabel("$I$", fontsize=get_font_size(i_fs)[1])
            ax.set_xscale('log')
            ax.set_yscale('log')
            idx = (self.data > 0)
            ax.errorbar(self.qgrid[idx], self.data[idx], self.err[idx], label=self.label)
            for ov in self.overlaps:
                ax.plot(ov['q_overlap'], ov['raw_data1'], "v")
                ax.plot(ov['q_overlap'], ov['raw_data2'], "^")

        d0 = copy.deepcopy(self)
        if len(dsets)==0:
            return d0

        if weighted:
            wt = []
            for d in dsets+[self]:
                w0 = np.sum(np.fabs(d.data[d.qgrid<qmax_for_weight]))
                if w0<=0:
                    raise Exception(f"weight for averaging <0: {w0}")
                wt.append(w0)
            wt = np.asarray(wt)
            wt /= wt.max()
            er2 = (self.err/wt[-1])**2
            d0.err = 1./er2
            d0.data = d0.data/er2
            if debug==True:
                print("weight factors: ", wt)

        for i in range(len(dsets)):
            d1 = dsets[i]
            if debug==True:
                print("%s " % d1.label, end=' ')
            if not (d0.qgrid == d1.qgrid).all():
                raise Exception("\n1D sets cannot be averaged: qgrid mismatch")

            d0.trans += d1.trans
            d0.trans_w += d1.trans_w
            d0.trans_e += d1.trans_e
            if weighted:
                er2 = (d1.err/wt[i])**2
                d0.err += 1/er2
                d0.data += d1.data/er2
            else:
                d0.data += d1.data
                d0.err += d1.err

            #if self.transMode == trans_mode.from_beam_center:
            #    d0.roi += d1.roi
            d0.comments += "# averaged with \n%s" % d1.comments.replace("# ", "## ")
            if plot_data:
                idx = (d1.data > 0)  # Remove Zeros on plot
                ax.errorbar(d1.qgrid[idx], d1.data[idx] * VOFFSET ** n, d1.err[idx] * VOFFSET ** n, label=d1.label)
            for i in range(len(d1.overlaps)):
                if plot_data:
                    ax.plot(d1.overlaps[i]['q_overlap'], d1.overlaps[i]['raw_data1'] * VOFFSET ** n, "v")
                    ax.plot(d1.overlaps[i]['q_overlap'], d1.overlaps[i]['raw_data2'] * VOFFSET ** n, "^")
                d0.overlaps[i]['raw_data1'] += d1.overlaps[i]['raw_data1']
                d0.overlaps[i]['raw_data2'] += d1.overlaps[i]['raw_data2']
            n += 1
            d0.label = common_name(d0.label, d1.label)

        d0.trans /= n
        d0.trans_w /= n
        d0.trans_e /= n
        if weighted:
            d0.data /= d0.err
            d0.err = 1./np.sqrt(d0.err)
        else:
            d0.data /= n
            d0.err /= (n*np.sqrt(n))   # should not be just sqrt(n), that would increase err after averaging
        #if self.transMode == trans_mode.from_beam_center:
        #    d0.roi /= n
        for ov in d0.overlaps:
            ov['raw_data1'] /= n
            ov['raw_data2'] /= n
        if debug==True:
            print("\naveraged set re-named to %s." % d0.label)

        if plot_data:
            # plot the averaged data over each individual curve
            for i in range(n):
                if i == 0:
                    idx = (d0.data > 0)  # Remove Zeros on plot
                    handles, labels = ax.get_legend_handles_labels()
                    lbl = "averaged" if "averaged" not in labels else ""
                    ax.plot(d0.qgrid[idx], d0.data[idx] * VOFFSET ** i, color="gray", lw=2, ls="--", label=lbl)
                else:
                    idx = (d0.data > 0)  # Remove Zeros on plot
                    ax.plot(d0.qgrid[idx], d0.data[idx] * VOFFSET ** i, color="gray", lw=2, ls="--")
            leg = ax.legend(loc='upper right', frameon=False)

            for t in leg.get_texts():
                t.set_fontsize(get_font_size(i_fs-2)[1])

        return d0

    def bkg_cor(self, dbak, sc_factor=1., plot_data=False, ax=None,
                inplace=False, check_overlap=False, show_eb=True, debug=False, fontsize='large'):
        """
        background subtraction
        """
        dset = None
        if inplace:
            dset = self
        else:
            dset = copy.deepcopy(self)
        i_fs = get_font_size(fontsize)[0]

        if debug==True:
            print("background subtraction: %s - %s" % (dset.label, dbak.label))
        if not (dbak.qgrid == dset.qgrid).all():
            print("background subtraction failed: qgrid mismatch")
            sys.exit()
        if dset.trans < 0 or dbak.trans <= 0:
            print("WARNING: trans value not assigned to data or background, assuming normalized intensity.")
            sc = 1.
        else:
            sc = dset.trans / dbak.trans

        # need to include raw data
        if plot_data:
            if ax is None:
                plt.figure()
                plt.subplots_adjust(bottom=0.15)
                ax = plt.gca()
            ax.set_xlabel("$q (\AA^{-1})$", fontsize=get_font_size(i_fs)[1])
            ax.set_ylabel("$I$", fontsize=get_font_size(i_fs)[1])
            ax.set_xscale('log')
            ax.set_yscale('log')
            ax.xaxis.set_tick_params(labelsize=get_font_size(i_fs-1)[1])
            ax.yaxis.set_tick_params(labelsize=get_font_size(i_fs-1)[1])
            idx = (dset.data > 0) & (dbak.data > 0)
            ax.plot(dset.qgrid[idx], dset.data[idx], label=self.label)
            ax.plot(dbak.qgrid[idx], dbak.data[idx], label=dbak.label)
            ax.plot(dbak.qgrid[idx], dbak.data[idx] * sc * sc_factor, label=dbak.label + ", scaled")

        if len(dset.overlaps) != len(dbak.overlaps):
            if check_overlap:
                raise Exception("Background subtraction failed: overlaps mismatch.")
        else:
            for i in range(len(dset.overlaps)):
                dset.overlaps[i]['raw_data1'] -= dbak.overlaps[i]['raw_data1'] * sc_factor * sc
                dset.overlaps[i]['raw_data2'] -= dbak.overlaps[i]['raw_data2'] * sc_factor * sc
                if plot_data:
                    ax.plot(dset.overlaps[i]['q_overlap'], dset.overlaps[i]['raw_data1'], "v")
                    ax.plot(dset.overlaps[i]['q_overlap'], dset.overlaps[i]['raw_data2'], "^")
        if plot_data:
            leg = ax.legend(loc='upper right', frameon=False)
            for t in leg.get_texts():
                t.set_fontsize(get_font_size(i_fs-2)[1])

        if debug==True:
            print("using scaling factor of %f" % (sc * sc_factor))
        dset.data -= dbak.data * sc * sc_factor
        dset.err += dbak.err * sc * sc_factor
        if plot_data:
            if show_eb:
                ax.errorbar(dset.qgrid, dset.data, dset.err)
            else:
                ax.plot(dset.qgrid, dset.data)

        dset.comments += "# background subtraction using the following set, scaled by %f (trans):\n" % sc
        if not sc_factor == 1.:
            dset.comments += "# with addtional scaling factor of: %f\n" % sc_factor
        dset.comments += dbak.comments.replace("# ", "## ")

        return dset

    def scale(self, sc):
        """
        scale the data by factor sc
        """
        if sc <= 0:
            print("scaling factor is non-positive: %f" % sc)
        self.data *= sc
        self.err *= sc
        self.trans *= sc
        if self.trans_w>0:
            self.trans_w *= sc
        if self.trans_e>0:
            self.trans_e *= sc
        self.comments += "# data is scaled by %f.\n" % sc
        if len(self.overlaps) != 0:
            for ov in self.overlaps:
                ov['raw_data1'] *= sc
                ov['raw_data2'] *= sc

        return self

    def merge(self, d1, qmax=-1, qmin=-1, fix_scale=-1, debug=False):
        """
        combine the data in self and d1
        scale d1 intensity to match self
        self and d1 should have the same qgrid

        if qmax or qmin <0
        simply keep the WAXS data that is beyond qmax for the SAXS data
        this is useful for utilizing WAXS to normalize intensity but keep SAXS data only
        """

        if debug==True:
            print("merging data: %s and %s ..." % (self.label, d1.label))
        if not (d1.qgrid == self.qgrid).all():
            print("merging data sets should have the same qgrid.")
            exit()

        # this gives the overlapping region
        idx = (self.data > 0) & (d1.data > 0)

        if len(self.qgrid[idx]) > 0:
            qmin0 = min(d1.qgrid[idx])
            qmax0 = max(self.qgrid[idx])
            # merge SAXS/WAXS based on intensity in the overlapping region
            if qmax0 < qmax:
                qmax = qmax0
            if qmin0 > qmin:
                qmin = qmin0
            idx = (self.qgrid > qmin) & (self.qgrid < qmax)
            # save the raw data in case needed, e.g. for ploting
            self.overlaps.append({'q_overlap': self.qgrid[idx],
                                  'raw_data1': self.data[idx],
                                  'raw_data2': d1.data[idx]})
        else:
            # no overlap
            # simply stack WAXS data to the high q end of SAXS data
            qmin = qmax = max(self.qgrid[self.data > 0])
            self.overlaps.append({'q_overlap': np.empty(0),
                                  'raw_data1': np.empty(0),
                                  'raw_data2': np.empty(0)})

        # idx = np.asarray([],dtype=int)

        if len(self.qgrid[idx])==0:
            if debug!='quiet':
                print("data sets are not overlapping in the given q range.")
            if fix_scale < 0:
                fix_scale = 1
                if debug!='quiet':
                    print("forcing fix_scale=1.")
        elif len(self.qgrid[idx]) < 5 and debug!='quiet':
            print("too few overlapping points: %d" % len(self.qgrid[idx]))

        if fix_scale > 0:
            # For a given experimental configuration, the intensity normlization
            # factor between the SAXS and WAXS should be well-defined. This factor
            # can be determined using scattering data with siginificant intensity
            # in the overlapping q-range and applied to all data collected in the
            # same configuration.
            sc = fix_scale
        else:
            sc = np.linalg.lstsq(np.asmatrix(self.data[idx]).T, np.asmatrix(d1.data[idx]).T)[0]
            sc = np.trace(sc)

        d1.data /= sc
        d1.err /= sc
        if len(self.qgrid[idx]) > 0:
            if debug==True:
                print("Scaled Overlaps by 1/%f" % sc)
            self.overlaps[-1]['raw_data2'] /= sc
            self.overlaps[-1]['sc'] = sc

        self.label = common_name(self.label, d1.label)
        if debug==True:
            print("set2 scaled by 1/%f" % sc)
            print("merged set re-named %s." % self.label)

        if len(self.qgrid[idx]) > 0:
            # averaging using 1/err^2 as weight, maximum likelihood estimator
            w1 = 1./(self.err[idx]*self.err[idx])
            w2 = 1./(d1.err[idx]*d1.err[idx])
            self.data[idx] = (w1*self.data[idx] + w2*d1.data[idx]) / (w1+w2)
            self.err[idx] = np.sqrt((w1/(w1+w2)*self.err[idx])**2 + (w2/(w1+w2)*d1.err[idx])**2)
        self.data[self.qgrid >= qmax] = d1.data[self.qgrid >= qmax]
        self.err[self.qgrid >= qmax] = d1.err[self.qgrid >= qmax]

        self.comments += "# merged with the following set by matching intensity within (%.4f, %.4f)," % (qmin, qmax)
        self.comments += " scaled by %f\n" % sc
        self.comments += d1.comments.replace("# ", "## ")

    def plot_Guinier(self, qs=0, qe=10, rg=15, fix_qe=False, scale_wabs=-1,
                     ax=None, no_plot=False, fontsize="large"):
        """ do Gunier plot, estimate Rg automatically
        qs specify the lower end of the q-range to perform the fit in
        rg is the optinal initial estimate
        if fix_qe==1, qe defined the end of the region to perform the fit
        """
        idx = (self.data > 0)
        i_fs = get_font_size(fontsize)[0]
        # print self.data

        scale = 1.0
        if scale_wabs>0:
            assert(self.trans_w>0)
            scale *= scale_wabs/self.trans_w

        if no_plot==False:
            if ax is None:
                ax = plt.gca()
            ax.set_xscale('linear')
            ax.set_yscale('log')
            ax.errorbar(self.qgrid[idx]**2, self.data[idx]*scale, self.err[idx]*scale)

        cnt = 0
        t = self.qgrid[self.data > 0][0]
        if qs < t: qs = t
        while cnt < 10:
            if (not fix_qe) and qe > 1.3/rg and 1.3/rg > qs+0.004: qe = 1.3/rg
            td = np.vstack((self.qgrid, self.data))
            td = td[:, td[0, :] >= qs]
            td = td[:, td[0, :] <= qe]
            td[0, :] = td[0, :] * td[0, :]
            td[1, :] = np.log(td[1, :])
            rg, i0 = np.polyfit(td[0, :], td[1, :], 1)
            i0 = np.exp(i0)
            if rg<0:
                rg = np.sqrt(-rg * 3.)
            else:
                rg = 1e-6   #
                print("likely strong inter-particle interaction ...")
                break
            cnt += 1
            # print i0, rg
        td[1, :] = i0 * np.exp(-td[0, :]*rg*rg/3.)
        n1 = len(self.qgrid[self.qgrid<qs])
        n2 = len(self.qgrid)-len(self.qgrid[self.qgrid>qe])
        fit_range = [n1,n2]
        i0*=scale

        if no_plot==False and rg>0.1:
            #ax.tick_params(axis='y', labelleft=False)
            ax.plot([td[0, 0], td[0, -1]], [td[1, 0]*scale, td[1, -1]*scale], "ro")
            ax.plot(self.qgrid**2, i0*np.exp(-(self.qgrid*rg)**2/3))
            ax.set_ylabel("$I$", fontsize=get_font_size(i_fs)[1])
            ax.set_xlabel("$q^2 (\AA^{-2})$", fontsize=get_font_size(i_fs)[1])
            ax.xaxis.set_tick_params(labelsize=get_font_size(i_fs-1)[1])
            ax.yaxis.set_tick_params(labelsize=get_font_size(i_fs-1)[1])
            # plt.subplots_adjust(bottom=0.15)
            ax.set_xlim(0, qe**2*1.2)
            #ax.autoscale_view(tight=True, scalex=False, scaley=True)
            if i0>0:
                ax.set_ylim(top=i0*2, bottom=i0*np.exp(-(qe*rg)**2/3)/2)
        # print "I0=%f, Rg=%f" % (i0,rg)
        return (i0, rg, fit_range) # include fit range to be more compatible with ATSAS

    def plot_pr(self, i0, rg, qmax=5., dmax=200., ax=None, fontsize='large'):
        """ calculate p(r) function
        use the given i0 and rg value to fill in the low q part of the gap in data
        truncate the high q end at qmax
        """
        i_fs = get_font_size(fontsize)[0]
        if ax is None:
            ax = plt.gca()
        ax.set_xscale('linear')
        ax.set_yscale('linear')

        if self.qgrid[-1] < qmax: qmax = self.qgrid[-1]
        tqgrid = np.arange(0, qmax, qmax / len(self.qgrid))
        tint = np.interp(tqgrid, self.qgrid, self.data)

        tint[tqgrid * rg < 1.] = i0 * np.exp(-(tqgrid[tqgrid * rg < 1.] * rg) ** 2 / 3.)
        # tint -= tint[-10:].sum()/10
        # Hanning window for reducing fringes in p(r)
        tw = np.hanning(2 * len(tqgrid) + 1)[len(tqgrid):-1]
        tint *= tw

        trgrid = np.arange(0, dmax, 1.)
        kern = np.asmatrix([[rj ** 2 * np.sinc(qi * rj / np.pi) for rj in trgrid] for qi in tqgrid])
        tt = np.asmatrix(tint * tqgrid ** 2).T
        tpr = np.reshape(np.array((kern.T * tt).T), len(trgrid))
        tpr /= tpr.sum()

        # plt.plot(tqgrid,tint,"g-")
        # tpr = np.fft.rfft(tint)
        # tx = range(len(tpr))
        ax.plot(trgrid, tpr, "g-")
        ax.set_xlabel("$r (\AA)$", fontsize=get_font_size(i_fs)[1])
        ax.set_ylabel("$P(r)$", fontsize=get_font_size(i_fs)[1])
        # plt.subplots_adjust(bottom=0.15)

    def save(self, fn, nz=True, scale_wabs=-1,
             save_comments=False, debug=False, footer=None):
        """
        should save all the relevant information, such as scaling, merging, averaging
        save data points with non-zero intensity only if nz==1
        """
        if scale_wabs>0:
            assert(self.trans_w>0)
            scale_wabs /= self.trans_w
            qidi = np.vstack((self.qgrid, self.data*scale_wabs, self.err*scale_wabs))
        else:
            qidi = np.vstack((self.qgrid, self.data, self.err))
        if nz:
            qidi = qidi[:, self.data != 0]
        if debug==True:
            print("saving file: %s, nz=%d" % (fn, nz))
        np.savetxt(fn, qidi.T, "%8.4f   %8.3e   %8.3e")
        if save_comments or footer is not None:
            ff = open(fn, "a")
            if save_comments:
                ff.write(self.comments)
                if scale_wabs>0:
                    ff.write("# converted to abs scale by applying a scaling factor of {scale_wabs:.2e}\n")
            elif footer is not None:
                ff.write(footer)
            ff.close()

    def plot(self, ax=None, scale=1., fontsize='large', scale_wabs=-1):
        i_fs = get_font_size(fontsize)[0]
        if ax is None:
            plt.figure()
            plt.subplots_adjust(bottom=0.15)
            ax = plt.gca()
        if scale_wabs>0:
            assert(self.trans_w>0)
            scale *= scale_wabs/self.trans_w
        ax.set_xlabel("$q (\AA^{-1})$", fontsize=get_font_size(i_fs)[1])
        ax.set_ylabel("$I$", fontsize=get_font_size(i_fs)[1])
        ax.set_xscale('log')
        ax.set_yscale('log')
        ax.errorbar(self.qgrid, self.data*scale, self.err*scale, label=self.label)
        for ov in self.overlaps:
            ax.plot(ov['q_overlap'], ov['raw_data1']*scale, "v")
            ax.plot(ov['q_overlap'], ov['raw_data2']*scale, "^")
        leg = ax.legend(loc='upper right', frameon=False)

        for t in leg.get_texts():
            t.set_fontsize(get_font_size(i_fs-2)[1])

def normalize(ds):
    return np.divide(ds.data, np.max(ds.data))

def calculate(ds0, ds1):
    diff_coef = np.sum(np.abs(np.subtract(ds0, ds1)))  # How different the datasets are
    return diff_coef

def filter_by_similarity(datasets, similarity_threshold=0.5, debug=False):

    if len(datasets)==1:
        return datasets,None

    number_of_cpus = os.cpu_count()
    number_of_datasets = len(datasets)
    combinations = list(it.combinations(range(number_of_datasets), 2))
    similarity_matrix = np.zeros((number_of_datasets, number_of_datasets), dtype=np.bool)
    np.fill_diagonal(similarity_matrix, 1)  # If we compare the dataset with itself the result will always be one.

    norm_data = collections.deque(map(normalize, datasets))
    with mp.Pool(number_of_cpus) as pool:
        differences = pool.starmap(calculate, [(norm_data[i], norm_data[j]) for i, j in combinations])

    # print("Differences: \n", differences)
    # print("Diff Norm: \n", np.divide(differences, np.max(differences)))
    # print("Combinations: \n", combinations)
    similarities = np.divide(differences, np.max(differences)) <= similarity_threshold

    idx = 0
    for c in combinations:
        similarity_matrix[c[0]][c[1]] = similarities[idx]
        similarity_matrix[c[1]][c[0]] = similarities[idx]
        idx += 1

    number_of_simil_per_column = np.sum(similarity_matrix, axis=0)

    # No valid candidate, return all the data
    if np.array_equal(number_of_simil_per_column, np.ones(number_of_datasets)):
        if debug is True:
            print("No dataset with similarity level below threshold. Returning everything.")
        return datasets, []

    best_datasets_column = np.argmax(number_of_simil_per_column)
    best_column = similarity_matrix[:, best_datasets_column]
    valid_entries = list(it.compress(datasets, best_column))
    invalid_entries = set(datasets) - set(valid_entries)
    # print("Similarity Matrix: \n", similarity_matrix)
    # print("Best Column: \n", best_column)
    return valid_entries, invalid_entries

def merge_detectors(fns, detectors, qgrid, reft=-1, plot_data=False, save_ave=False, save_merged=False, ax=None, qmax=-1, qmin=-1, fix_scale=1, debug=False, transMode=trans_mode.from_waxs, trans=-1):
    """
    fns: filename, without the _SAXS/_WAXS suffix
    fix_scale is now default to 1
    implicitly assume that all detectors have the same qgrid
    """
    ss = []
    t0 = time.time()
    for fn in fns:
        s0 = Data1d()
        d_tot = np.zeros(qgrid.shape)
        d_max = np.zeros(qgrid.shape)
        d_min = np.zeros(qgrid.shape)+1.e32
        e_tot = np.zeros(qgrid.shape)
        c_tot = np.zeros(qgrid.shape)
        label = None
        comments = ""
        #t1 = time.time()
        for det in detectors:
            #t2 = time.time()
            # revised 2017mar10
            # old conversion: fn+det.extension gives the complete filename, the extension looks like this: "_SAXS.cbf"
            # this is a problem when the detector collect multiple images per trigger
            # new comvention: fn is a template, e.g. '/GPFS/xf16id/exp_path/301525/301016/temp1_000002%s_00001.cbf',
            # and the extension looks like this: "_SAXS"
            if "%s" in fn:
                fn1 = fn % det.extension
            else:
                fn1 = fn + det.extension
            if debug==True:
                print(fn, det.extension, fn1)

            s0.load_from_2D(fn1, det.exp_para, qgrid, det.pre_process,
                            save_ave=save_ave, debug=debug)

            if save_ave:
                s0.save(fn1 + ".ave", debug=debug)

            if det.fix_scale is not None:
                fix_scale = det.fix_scale
                s0.scale(1./fix_scale)

            # empty part of the data is nan
            idx = ~np.isnan(s0.data)
            d_tot[idx] += s0.data[idx]
            e_tot[idx] += s0.err[idx]
            c_tot[idx] += 1

            idx1 = (np.ma.fix_invalid(s0.data, fill_value=-1)>d_max).data
            d_max[idx1] = s0.data[idx1]
            idx2 = (np.ma.fix_invalid(s0.data, fill_value=1e32)<d_min).data
            d_min[idx2] = s0.data[idx2]

            comments += s0.comments
            if label is None:
                label = s0.label
            else:
                label = common_name(label, s0.label)

        s0.data = d_tot
        s0.err = e_tot
        idx = (c_tot>1)
        s0.overlaps.append({'q_overlap': qgrid[idx],
                             'raw_data1': d_max[idx],
                             'raw_data2': d_min[idx]})
        s0.data[idx] /= c_tot[idx]
        s0.err[idx] /= c_tot[idx]
        s0.transMode = transMode
        s0.set_trans(transMode, trans=trans, ref_trans=reft, debug=debug)
        s0.label = label
        s0.comments = comments # .replace("# ", "## ")

        if save_merged:
            s0.save(s0.label + ".dd", debug=debug)
        ss.append(s0)

    return ss

def average(fns, detectors, qgrid, reft=-1, plot_data=False, save1d=0, ax=None, qmax=-1, qmin=-1, fix_scale=-1,
            filter_datasets=True, similarity_threshold=0.5, debug=False):
    """
    fns: filename, without the _SAXS/_WAXS suffix
    save1d:  0 = do not save 1d data; 1 = save only the averaged data *.ddd;
             2 = merged data before average as well, *.dd; 3 = also save the data before merge, *.ave
    """

    save_dd = False
    save_ave = False
    save_ddd = False
    if save1d>0:
        save_ddd = True
    if save1d>1:
        save_dd = True
    if save1d>2:
        save_ave = True

    t0 = time.time()
    ss = merge_detectors(fns, detectors, qgrid, reft, plot_data, save_ave, save_dd,
                         ax, qmax, qmin, fix_scale, debug=debug)
    t1 = time.time()
    if filter_datasets:
        ss, invalids = filter_by_similarity(ss, similarity_threshold=similarity_threshold, debug=debug)
        # TODO: Insert warning/exception when the number of datasets discarded is high.
        # TODO: Define the % of discarded to result in a error.
        if debug!='quiet':
            print("Selected Datasets: ")
            for s in ss:
                print(s.label)

        if len(invalids) > 0 and debug!='quiet':
            print("The following datasets where discarded due to similarity level below the threshold: ",
                  similarity_threshold)
            for inv in invalids:
                print(inv.label)
    t2 = time.time()
    if len(ss) > 0:
        ss[0] = ss[0].avg(ss[1:], plot_data, ax=ax, debug=debug)
    t3 = time.time()
    if save_ddd:
        ss[0].save(ss[0].label + ".ddd", debug=debug)
    t4 = time.time()
    if debug==True:
        print('Time to Merge: ', t1-t0)
        print('Time to Filter: ', t2-t1)
        print('Time to Average: ', t3-t2)
        print('Time to Save Data: ', t4-t3)
    return ss[0]

def process(sfns, bfns, detectors, qgrid, qmax=-1, qmin=-1,
            reft=-1, save1d=False, conc=0., plot_data=True, fix_scale=-1,
            filter_datasets=True, similarity_threshold=0.5, debug=False):
    vfrac = 0.001 * conc / PROTEIN_WATER_DENSITY_RATIO

    sample_axis = None
    buffer_axis = None
    bkg_cor_axis = None

    if plot_data:
        plt.figure()
        sample_axis = plt.subplot2grid((2, 2), (0, 0), title="Sample Data")
        buffer_axis = plt.subplot2grid((2, 2), (0, 1), title="Buffer Data")
        bkg_cor_axis = plt.subplot2grid((2, 2), (1, 0), colspan=2, title="Background Correction")

    # TODO: Run the next two lines in parallel
    args_sample = (sfns, detectors, qgrid, reft, plot_data, save1d, sample_axis, qmax,
                   qmin, fix_scale, filter_datasets, similarity_threshold)
    args_buffer = (bfns, detectors, qgrid, reft, plot_data, save1d, buffer_axis, qmax,
                   qmin, fix_scale, filter_datasets, similarity_threshold)

    # In order for the Pool to work we need to make the ExpPara something other than a SwigObject.
    # Maybe it is time to change for something else
    # with mp.Pool(1) as pool:
    #     result = pool.starmap(average, (args_sample, args_buffer))
    #
    # ds = result[0]
    # db = result[1]

    ds = average(*args_sample, debug=debug)
    db = average(*args_buffer, debug=debug)

    ds.bkg_cor(db, 1.0 - vfrac, plot_data=plot_data, ax=bkg_cor_axis, inplace=True, debug=debug)

    return ds

def analyze(d1, qstart, qend, fix_qe, qcutoff, dmax):
    plt.figure(figsize=(14, 5.5))
    plt.subplot(121)
    I0,Rg,_ = d1.plot_Guinier(qs=qstart, qe=qend, fix_qe=fix_qe)

    print("I0=%f, Rg=%f" % (I0, Rg))

    plt.subplot(122)
    d1.plot_pr(I0, Rg, qmax=1.2, dmax=dmax)
    plt.subplots_adjust(bottom=0.15, wspace=0.25)

def estimate_scaling_factor(d1s, d1b,
                            q_min=0.5, q_max=3.5, smoothing_width=5, prec=4, s_thresh=1,
                            plot_data=False, ax=None, debug=False):
    """ Estimate the scaling factor needed to subtract buffer scattering d1b
        from sample scattering d1s, d1s/d1b should be instances of Data1d

        This function iteratively vary the scaling factor, up to the specified
        precision (number of digits after the decimal point), and optinally smooth
        the data before the calcualtion. Several critiera for over sub-traction are
        used to stop the interation:
        (1) non-zero value in the subtracted result
        (2) minimum of the subtracted result fall below 1/q
        (3) dynamic range of the value of the result, as measured by the span or std
            deviation of the log, becomes too high

        only test the data above q_min, which should not exceed 1.0
    """
    idx = (d1s.qgrid>q_min) & (d1s.qgrid<q_max) & (d1s.data>0) & (d1b.data>0)
    md1s = d1s.data[idx]
    md1b = d1b.data[idx]
    try:
        md1sm = smooth(md1s, half_window_len=smoothing_width)
        md1bm = smooth(md1b, half_window_len=smoothing_width)
    except:
        return 1.0      # problems like incomplete data can cause smooth() to fail, no point to continue
    mq = d1s.qgrid[idx]
    sc = 0.9
    prec0 = 2
    std0 = np.log(md1s-md1b*sc).std()
    n,bins = np.histogram((md1s-md1b*sc)*mq)
    sp0 = bins[-1]-bins[0]

    if debug:
        print("# sc,   mean,   std,   span,   qmin")
    while prec0<=prec:
        sc1 = sc
        while np.all(md1s-sc1*md1b>0):
            # data should have finite dynamic range
            td = np.log(md1s-md1b*sc1)
            n,bins = np.histogram(td)
            sp1 = bins[-1]-bins[0]
            std1 = td.std()
            #if std1>s_thresh*std0:
            #    break
            #elif std1<std0:
            #    std0 = std1
            if sp1>sp0+s_thresh:
                if debug:
                    print(f"span exceeded threshold: {sp1}, {sp0}+{s_thresh}")
                break
            # assume that data*q should have a lower bound
            # this should work better on the smoothed data
            td1 = (md1sm-md1bm*sc1)*mq
            q_Imin = mq[td1.argmin()]
            if q_Imin>1.7:  # under the water peak, indication of over-subtraction
                if debug:
                    print(f"qxI min under water peak: q={q_Imin:.3f}")
                break
            sc = sc1
            if debug:
                print(f"{sc1:.5f}, {td.mean():.3f}, {std1:.3f}, {sp1:3f}, {q_Imin:.3f}")
            sc1 = sc+np.power(10., -prec0)
        prec0 += 1

    if plot_data:
        if ax is None:
            plt.figure()
            ax = plt.gca()
        plt.semilogy(d1s.qgrid, d1s.data)
        plt.semilogy(d1b.qgrid, d1b.data)
        plt.errorbar(d1s.qgrid, (d1s.data-sc*d1b.data), d1s.err*1.414)

    return sc


'''
@Author: dengzaiyong
@Date: 2021-08-21 15:16:08
@LastEditTime: 2021-08-27 19:37:08
@LastEditors: dengzaiyong
@Desciption: 使用Faiss训练hnsw模型
@FilePath: /JDQA/retrieval/hnsw_faiss.py
'''

import os
import time
import numpy as np
import pandas as pd
from gensim.models import KeyedVectors
import faiss
import config
from preprocessor import clean
from utils.tools import create_logger

logger = create_logger(config.root_path + '/logs/hnsw_faiss.log')

def sentence_embedding(sentence, w2v_model):
    '''
    通过词向量均值的方式生成句向量
    sentence: 待生成句向量的句子
    w2v_model: word2vec模型
    return: 句子中所有词向量的均值
    '''
    embedding = []
    for word in clean(sentence).split():
        if word not in w2v_model.wv.index_to_key:
            embedding.append(np.random.randn(1, config.embed_dim))
        else:
            embedding.append(w2v_model.wv.get_vector(word))

    # 所有词向量的均值为句向量
    return np.mean(np.array(embedding), axis=0).reshape(1, -1)

class HNSW(object):
    def __init__(self,
                 w2v_path,                      # word2vec模型路径
                 ef=config.ef_construction,     # 搜索时保存最近邻的动态列表大小
                 M=config.M,                    # 节点的邻结点的数量
                 model_path=None,               # hnsw模型保存路径
                 data_path=None):               # 数据文件路径

        self.w2v_model = KeyedVectors.load(w2v_path)
        self.data = self.load_data(data_path)

        # 加载hnsw模型
        if model_path and os.path.exists(model_path):
            self.index = self.load_hnsw(model_path)
        # 训练hnsw模型
        elif data_path:
            self.index = self.build_hnsw(model_path, ef=ef, m=M)
        else:
            logger.error('No existing model and no building data provided.')

    def load_data(self, data_path):
        '''
        读取数据，并生成句向量
        data_path：问答pair数据所在路径
        return: 包含句向量的dataframe
        '''
        data = pd.read_csv(data_path)

        # 生成custom每条记录的句向量
        data['custom_vec'] = data['custom'].apply(
            lambda x: sentence_embedding(x, self.w2v_model))

        # 确保句向量的维度为300
        data['custom_vec'] = data['custom_vec'].apply(
            lambda x: x[0][0] if x.shape[1] != config.embed_dim else x)
        data = data.dropna()

        return data

    def evaluate(self, vecs):
        '''
        验证模型
        '''
        logger.info('Evaluating hnsw model')
        nq, d = vecs.shape
        t0 = time.time

        # 找top1个相似的
        D, I = self.index.search(vecs, 1)
        t1 = time.time

        missing_rate = (I == -1).sum() / float(nq)
        recall_at_1 = (I == np.arange(nq)).sum() / float(nq)
        print("\t %7.3f ms per query, R@1 %.4f, missing rate %.4f" % (
            (t1 - t0) * 1000.0 / nq, recall_at_1, missing_rate))

    def build_hnsw(self, to_file, ef=2000, m=64):
        """
        训练hnsw模型
        """
        logger.info('building hnsw index')

        # 所有的句向量拼接
        vecs = np.stack(self.data['custom_vec'].values).reshape(-1, config.embed_dim)
        vecs = vecs.astype('float32')

        dim = self.w2v_model.vector_size

         # 构建索引
        index = faiss.IndexHNSWFlat(dim, m)

        # 使用单个GPU资源
        res = faiss.StandardGpuResources()

        faiss.index_cpu_to_gpu(res, 0, index)
        index.hnsw.ef_construction = ef
        index.verbose = True
        index.add(vecs)

        # 保存hnsw模型
        faiss.write_index(index, to_file)

        return index

    def load_hnsw(self, model_path):
        logger.info(f"Loading hnsw from {model_path}")
        hnsw = faiss.read_index(model_path)
        return hnsw

    def search(self, text, k=5):
        """
        通过hnsw检索topk
        """
        print(f'Searching for {text}.')

        # 转换句向量
        test_vec = sentence_embedding(clean(text), self.w2v_model)
        test_vec = test_vec.astype('float32')

        # 搜索相似度最高的k个句向量
        D, I = self.index.search(test_vec, k)

        df = pd.concat((self.data.iloc[I[0]]['custom'].reset_index(),
                        self.data.iloc[I[0]]['assistance'].reset_index(drop=True),
                        pd.DataFrame(D.reshape(-1, 1), columns=['q_distance'])),
                        axis=1)

        return df

if __name__ == '__main__':
    hnsw = HNSW(config.w2v_path,
            config.ef_construction,
            config.M,
            config.fassi_path,
            config.train_path)

    text = '我要转人工'
    print(hnsw.search(text, k=10))

    # 验证模型
    eval_vecs = np.stack(hnsw.data['custom_vec'].values).reshape(-1, config.embed_dim)
    eval_vecs.astype('float32')
    hnsw.evaluate(eval_vecs[:1000])


"""
Copyright 2017 Steven Diamond

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
from collections import namedtuple

import cvxpy.interface as intf
import cvxpy.settings as s
import cvxpy.lin_ops.lin_utils as lu
import numpy as np
from cvxpy.problems.solvers.solver import Solver
from scipy.sparse import dok_matrix

# Values used to distinguish between linear and quadratic constraints.
_LIN, _QUAD = 0, 1
# For internal bookkeeping, we have to separate linear indices from
# quadratic indices. The "cpx_constrs" member of the results_dict will
# contain namedtuples of (constr_type, index) where constr_type is either
# _LIN or _QUAD.
_CpxConstr = namedtuple("_CpxConstr", ["constr_type", "index"])

class CPLEX(Solver):
    """An interface for the CPLEX solver.
    """

    # Solver capabilities.
    LP_CAPABLE = True
    SOCP_CAPABLE = True
    SDP_CAPABLE = False
    EXP_CAPABLE = False
    MIP_CAPABLE = True

    def name(self):
        """The name of the solver.
        """
        return s.CPLEX

    def import_solver(self):
        """Imports the solver.
        """
        import cplex
        cplex  # For flake8

    def matrix_intf(self):
        """The interface for matrices passed to the solver.
        """
        return intf.DEFAULT_SPARSE_INTF

    def vec_intf(self):
        """The interface for vectors passed to the solver.
        """
        return intf.DEFAULT_INTF

    def split_constr(self, constr_map):
        """Extracts the equality, inequality, and nonlinear constraints.

        Parameters
        ----------
        constr_map : dict
            A dict of the canonicalized constraints.

        Returns
        -------
        tuple
            (eq_constr, ineq_constr, nonlin_constr)
        """
        return (constr_map[s.EQ] + constr_map[s.LEQ], [], [])

    @staticmethod
    def _param_in_constr(constraints):
        """Do any of the constraints contain parameters?
        """
        for constr in constraints:
            if len(lu.get_expr_params(constr.expr)) > 0:
                return True
        return False

    def solve(self, objective, constraints, cached_data,
              warm_start, verbose, solver_opts):
        """Returns the result of the call to the solver.

        Parameters
        ----------
        objective : LinOp
            The canonicalized objective.
        constraints : list
            The list of canonicalized cosntraints.
        cached_data : dict
            A map of solver name to cached problem data.
        warm_start : bool
            Not used.
        verbose : bool
            Should the solver print output?
        solver_opts : dict
            Additional arguments for the solver.
            'cplex_params' - a dictionary where the key-value pairs are
                             composed of parameter names and parameter
                             values.
            'cplex_filename' - A string specifying the filename to which
                               the problem will be written.

        Returns
        -------
        tuple
            (status, optimal value, primal, equality dual, inequality dual)
        """
        import cplex

        # Get problem data
        data = self.get_problem_data(objective, constraints, cached_data)

        c = data[s.C]
        b = data[s.B]
        A = dok_matrix(data[s.A])
        # Save the dok_matrix.
        data[s.A] = A
        data[s.BOOL_IDX] = solver_opts[s.BOOL_IDX]
        data[s.INT_IDX] = solver_opts[s.INT_IDX]

        n = c.shape[0]

        solver_cache = cached_data[self.name()]

        # TODO: warmstart with SOC constraints.
        if warm_start and solver_cache.prev_result is not None \
           and len(data[s.DIMS][s.SOC_DIM]) == 0:
            model = solver_cache.prev_result["model"]
            variables = solver_cache.prev_result["variables"]
            # cpx_constrs contains CpxConstr namedtuples (see above).
            cpx_constrs = solver_cache.prev_result["cpx_constrs"]
            c_prev = solver_cache.prev_result["c"]
            A_prev = solver_cache.prev_result["A"]
            b_prev = solver_cache.prev_result["b"]

            # If there is a parameter in the objective, it may have changed.
            if len(lu.get_expr_params(objective)) > 0:
                c_diff = c - c_prev

                I_unique = list(set(np.where(c_diff)[0]))

                for i in I_unique:
                    model.objective.set_linear(variables[i], c[i])
            else:
                # Stay consistent with CPLEX's representation of the problem
                c = c_prev

            # Get equality and inequality constraints.
            sym_data = self.get_sym_data(objective, constraints, cached_data)
            all_constrs, _, _ = self.split_constr(sym_data.constr_map)

            # If there is a parameter in the constraints,
            # A or b may have changed.
            if self._param_in_constr(all_constrs):
                A_diff = dok_matrix(A - A_prev)
                b_diff = b - b_prev

                # Figure out which rows of A and elements of b have changed
                try:
                    idxs, _ = zip(*[x for x in A_diff.keys()])
                except ValueError:
                    idxs = []
                I_unique = list(set(idxs) | set(np.where(b_diff)[0]))

                # Update locations which have changed
                csr = A.tocsr()
                for i in I_unique:
                    # To update a constraint, we first disable the old
                    # constraint and then add a new constraint with the
                    # modifications. This way we don't have to worry
                    # about indices needing to shift. The old constraint
                    # is disabled by setting all coefficients and the rhs
                    # to zero.
                    #
                    # NOTE: This can change the relative order of the
                    # constraints, which can result in performance
                    # variability!

                    # Disable the old constraint if it exists.
                    assert cpx_constrs[i].index is not None
                    assert cpx_constrs[i].constr_type == _LIN
                    idx = cpx_constrs[i].index
                    tmp = model.linear_constraints.get_rows(idx)
                    model.linear_constraints.set_linear_components(
                        idx,
                        cplex.SparsePair(ind=tmp.ind, val=[0.0]*len(tmp.ind)))
                    model.linear_constraints.set_rhs(idx, 0.0)

                    # Add new constraint

                    ind = [variables[x] for x in csr[i].indices]
                    val = [x for x in csr[i].data]
                    if i < data[s.DIMS][s.EQ_DIM]:
                        ctype = "E"
                    else:
                        assert data[s.DIMS][s.EQ_DIM] <= i \
                            < data[s.DIMS][s.EQ_DIM] + data[s.DIMS][s.LEQ_DIM]
                        ctype = "L"
                    new_idx = list(model.linear_constraints.add(
                        lin_expr=[cplex.SparsePair(ind=ind, val=val)],
                        senses=ctype,
                        rhs=[b[i]]))[0]
                    cpx_constrs[i] = _CpxConstr(_LIN, new_idx)

            else:
                # Stay consistent with CPLEX's representation of the problem
                A = A_prev
                b = b_prev

        else:
            model = cplex.Cplex()
            variables = []
            # cpx_constrs will contain CpxConstr namedtuples (see above).
            cpx_constrs = []
            vtype = []
            if self.is_mip(data):
                for i in range(n):
                    # Set variable type.
                    if i in data[s.BOOL_IDX]:
                        vtype.append('B')
                    elif i in data[s.INT_IDX]:
                        vtype.append('I')
                    else:
                        vtype.append('C')
            else:
                # If we specify types (even with 'C'), then the problem will
                # be interpreted as a MIP. Leaving vtype as an empty list
                # here, will ensure that the problem type remains an LP.
                pass
            # Add the variables in a batch
            variables = list(model.variables.add(
                obj=[c[i] for i in range(n)],
                lb=[-cplex.infinity]*n,  # default LB is 0
                ub=[cplex.infinity]*n,
                types="".join(vtype),
                names=["x_%d" % i for i in range(n)]))

            # Add equality constraints
            cpx_constrs += [_CpxConstr(_LIN, x)
                            for x in self.add_model_lin_constr(
                                    model, variables,
                                    range(data[s.DIMS][s.EQ_DIM]),
                                    'E', A, b)]

            # Add inequality (<=) constraints
            leq_start = data[s.DIMS][s.EQ_DIM]
            leq_end = data[s.DIMS][s.EQ_DIM] + data[s.DIMS][s.LEQ_DIM]
            cpx_constrs += [_CpxConstr(_LIN, x)
                            for x in self.add_model_lin_constr(
                                    model, variables,
                                    range(leq_start, leq_end),
                                    'L', A, b)]

            # Add SOC constraints
            soc_start = leq_end
            for constr_len in data[s.DIMS][s.SOC_DIM]:
                soc_end = soc_start + constr_len
                soc_constr, new_leq, new_vars = self.add_model_soc_constr(
                    model, variables, range(soc_start, soc_end), A, b)
                cpx_constrs.append(_CpxConstr(_QUAD, soc_constr))
                cpx_constrs += [_CpxConstr(_LIN, x) for x in new_leq]
                variables += new_vars
                soc_start += constr_len

        # Set verbosity
        if not verbose:
            model.set_results_stream(None)
            model.set_warning_stream(None)
            model.set_error_stream(None)
            model.set_log_stream(None)
        else:
            # By default the output will be sent to stdout.
            pass

        # TODO: user option to not compute duals.
        model.parameters.preprocessing.qcpduals.set(
            model.parameters.preprocessing.qcpduals.values.force)

        # TODO: Parameter support is functional, but perhaps not ideal.
        # The user must pass parameter names as used in the CPLEX Python
        # API, and raw values (i.e., no enum support).
        kwargs = sorted(solver_opts.keys())
        if "cplex_params" in kwargs:
            for param, value in solver_opts["cplex_params"].items():
                try:
                    eval("model.parameters.{0}.set({1})".format(param, value))
                except AttributeError:
                    raise ValueError(
                        "invalid CPLEX parameter, value pair ({0}, {1})".format(
                            param, value))
            kwargs.remove("cplex_params")
        if "cplex_filename" in kwargs:
            filename = solver_opts["cplex_filename"]
            if filename:
                model.write(filename)
            kwargs.remove("cplex_filename")
        if s.BOOL_IDX in kwargs:
            kwargs.remove(s.BOOL_IDX)
        if s.INT_IDX in kwargs:
            kwargs.remove(s.INT_IDX)
        if kwargs:
            raise ValueError("invalid keyword-argument '{0}'".format(kwargs[0]))

        results_dict = {}
        start_time = model.get_time()
        solve_time = -1
        try:
            model.solve()
            solve_time = model.get_time() - start_time
            results_dict["primal objective"] = model.solution.get_objective_value()
            results_dict["x"] = np.array(model.solution.get_values(variables))
            results_dict["status"] = self._get_status(model)

            # Only add duals if not a MIP.
            if not self.is_mip(data):
                vals = []
                for con in cpx_constrs:
                    assert con.index is not None
                    if con.constr_type == _LIN:
                        vals.append(model.solution.get_dual_values(con.index))
                    else:
                        assert con.constr_type == _QUAD
                        # Quadratic constraints not queried directly.
                        vals.append(0.0)
                results_dict["y"] = -np.array(vals)
        except Exception:
            if solve_time < 0.0:
                solve_time = model.get_time() - start_time
            results_dict["status"] = s.SOLVER_ERROR

        results_dict["model"] = model
        results_dict["variables"] = variables
        results_dict["cpx_constrs"] = cpx_constrs
        results_dict[s.SOLVE_TIME] = solve_time

        return self.format_results(results_dict, data, cached_data)

    def _handle_solve_status(self, model, solstat):
        """Map CPLEX MIP solution status codes to non-MIP status codes."""
        status = model.solution.status
        if solstat == status.MIP_optimal:
            return status.optimal
        elif solstat == status.MIP_infeasible:
            return status.infeasible
        elif solstat in (status.MIP_time_limit_feasible,
                         status.MIP_time_limit_infeasible):
            return status.abort_time_limit
        elif solstat in (status.MIP_dettime_limit_feasible,
                         status.MIP_dettime_limit_infeasible):
            return status.abort_dettime_limit
        elif solstat in (status.MIP_abort_feasible,
                         status.MIP_abort_infeasible):
            return status.abort_user
        elif solstat == status.MIP_optimal_infeasible:
            return status.optimal_infeasible
        elif solstat == status.MIP_infeasible_or_unbounded:
            return status.infeasible_or_unbounded
        elif solstat in (status.MIP_unbounded,
                         status.MIP_benders_master_unbounded,
                         status.benders_master_unbounded):
            return status.unbounded
        elif solstat in (status.feasible_relaxed_sum,
                         status.MIP_feasible_relaxed_sum,
                         status.optimal_relaxed_sum,
                         status.MIP_optimal_relaxed_sum,
                         status.feasible_relaxed_inf,
                         status.MIP_feasible_relaxed_inf,
                         status.optimal_relaxed_inf,
                         status.MIP_optimal_relaxed_inf,
                         status.feasible_relaxed_quad,
                         status.MIP_feasible_relaxed_quad,
                         status.optimal_relaxed_quad,
                         status.MIP_optimal_relaxed_quad):
            raise AssertionError(
                "feasopt status encountered: {0}".format(solstat))
        elif solstat in (status.conflict_feasible,
                         status.conflict_minimal,
                         status.conflict_abort_contradiction,
                         status.conflict_abort_time_limit,
                         status.conflict_abort_dettime_limit,
                         status.conflict_abort_iteration_limit,
                         status.conflict_abort_node_limit,
                         status.conflict_abort_obj_limit,
                         status.conflict_abort_memory_limit,
                         status.conflict_abort_user):
            raise AssertionError(
                "conflict refiner status encountered: {0}".format(solstat))
        elif solstat == status.relaxation_unbounded:
            return status.relaxation_unbounded
        elif solstat in (status.feasible,
                         status.MIP_feasible):
            return status.feasible
        elif solstat == status.benders_num_best:
            return status.num_best
        else:
            return solstat

    def _get_status(self, model):
        """Map CPLEX status to CPXPY status."""
        pfeas = model.solution.is_primal_feasible()
        # NOTE: dfeas is always false for a MIP.
        dfeas = model.solution.is_dual_feasible()
        status = model.solution.status
        solstat = self._handle_solve_status(model, model.solution.get_status())
        if solstat in (status.node_limit_infeasible,
                       status.fail_infeasible,
                       status.mem_limit_infeasible,
                       status.fail_infeasible_no_tree,
                       status.num_best):
            return s.SOLVER_ERROR
        elif solstat in (status.abort_user,
                         status.abort_iteration_limit,
                         status.abort_time_limit,
                         status.abort_dettime_limit,
                         status.abort_obj_limit,
                         status.abort_primal_obj_limit,
                         status.abort_dual_obj_limit,
                         status.abort_relaxed,
                         status.first_order):
            if pfeas:
                return s.OPTIMAL_INACCURATE
            else:
                return s.SOLVER_ERROR
        elif solstat in (status.node_limit_feasible,
                         status.solution_limit,
                         status.populate_solution_limit,
                         status.fail_feasible,
                         status.mem_limit_feasible,
                         status.fail_feasible_no_tree,
                         status.feasible):
            if dfeas:
                return s.OPTIMAL
            else:
                return s.OPTIMAL_INACCURATE
        elif solstat in (status.optimal,
                         status.optimal_tolerance,
                         status.optimal_infeasible,
                         status.optimal_populated,
                         status.optimal_populated_tolerance):
            return s.OPTIMAL
        elif solstat in (status.infeasible,
                         status.optimal_relaxed_sum,
                         status.optimal_relaxed_inf,
                         status.optimal_relaxed_quad):
            return s.INFEASIBLE
        elif solstat in (status.feasible_relaxed_quad,
                         status.feasible_relaxed_inf,
                         status.feasible_relaxed_sum):
            return s.SOLVER_ERROR
        elif solstat == status.infeasible_or_unbounded:
            return s.INFEASIBLE
        elif solstat == status.unbounded:
            return s.UNBOUNDED
        else:
            return s.SOLVER_ERROR

    def add_model_lin_constr(self, model, variables,
                             rows, ctype, mat, vec):
        """Adds EQ/LEQ constraints to the model using the data from mat and vec.

        Parameters
        ----------
        model : CPLEX model
            The problem model.
        variables : list
            The problem variables.
        rows : range
            The rows to be constrained.
        ctype : CPLEX constraint type
            The type of constraint.
        mat : SciPy COO matrix
            The matrix representing the constraints.
        vec : NDArray
            The RHS part of the constraints.

        Returns
        -------
        list
            A list of new linear constraint indices.
        """
        constr, lin_expr, rhs = [], [], []
        csr = mat.tocsr()
        for i in rows:
            ind = [variables[x] for x in csr[i].indices]
            val = [x for x in csr[i].data]
            lin_expr.append([ind, val])
            rhs.append(vec[i])
        # For better performance, we add the contraints in a batch.
        if lin_expr:
            assert len(lin_expr) == len(rhs)
            constr.extend(list(
                model.linear_constraints.add(
                    lin_expr=lin_expr,
                    senses=ctype * len(lin_expr),
                    rhs=rhs)))
        return constr

    def add_model_soc_constr(self, model, variables,
                             rows, mat, vec):
        """Adds SOC constraint to the model using the data from mat and vec.

        Parameters
        ----------
        model : CPLEX model
            The problem model.
        variables : list
            The problem variables.
        rows : range
            The rows to be constrained.
        mat : SciPy COO matrix
            The matrix representing the constraints.
        vec : NDArray
            The RHS part of the constraints.

        Returns
        -------
        tuple
            A tuple of (a new quadratic constraint index, a list of new
            supporting linear constr indices, and a list of new
            supporting variable indices).
        """
        import cplex
        # Assume first expression (i.e. t) is nonzero.
        lin_expr_list, soc_vars, lin_rhs = [], [], []
        csr = mat.tocsr()
        for i in rows:
            ind = [variables[x] for x in csr[i].indices]
            val = [x for x in csr[i].data]
            # Ignore empty constraints.
            if ind:
                lin_expr_list.append((ind, val))
                lin_rhs.append(vec[i])
            else:
                lin_expr_list.append(None)
                lin_rhs.append(0.0)

        # Make a variable and equality constraint for each term.
        soc_vars, is_first = [], True
        for i in rows:
            if is_first:
                lb = [0.0]
                names = ["soc_t_%d" % i]
                is_first = False
            else:
                lb = [-cplex.infinity]
                names = ["soc_x_%d" % i]
            soc_vars.extend(list(model.variables.add(
                obj=[0],
                lb=lb,
                ub=[cplex.infinity],
                types="",
                names=names)))

        new_lin_constrs = []
        for i, expr in enumerate(lin_expr_list):
            if expr is None:
                ind = [soc_vars[i]]
                val = [1.0]
            else:
                ind, val = expr
                ind.append(soc_vars[i])
                val.append(1.0)
            new_lin_constrs.extend(list(
                model.linear_constraints.add(
                    lin_expr=[cplex.SparsePair(ind=ind, val=val)],
                    senses="E",
                    rhs=[lin_rhs[i]])))

        assert len(soc_vars) > 0
        qconstr = model.quadratic_constraints.add(
            lin_expr=cplex.SparsePair(ind=[], val=[]),
            quad_expr=cplex.SparseTriple(
                ind1=soc_vars,
                ind2=soc_vars,
                val=[-1.0] + [1.0] * (len(soc_vars) - 1)),
            sense="L",
            rhs=0.0,
            name="")
        return (qconstr, new_lin_constrs, soc_vars)

    def format_results(self, results_dict, data, cached_data):
        """Converts the solver output into standard form.

        Parameters
        ----------
        results_dict : dict
            The solver output.
        data : dict
            Information about the problem.
        cached_data : dict
            A map of solver name to cached problem data.

        Returns
        -------
        dict
            The solver output in standard form.
        """
        dims = data[s.DIMS]
        if results_dict["status"] != s.SOLVER_ERROR:
            solver_cache = cached_data[self.name()]
            solver_cache.prev_result = {
                "model": results_dict["model"],
                "variables": results_dict["variables"],
                "cpx_constrs": results_dict["cpx_constrs"],
                "c": data[s.C],
                "A": data[s.A],
                "b": data[s.B],
            }
        new_results = {}
        new_results[s.STATUS] = results_dict['status']
        new_results[s.SOLVE_TIME] = results_dict[s.SOLVE_TIME]
        if new_results[s.STATUS] in s.SOLUTION_PRESENT:
            primal_val = results_dict['primal objective']
            new_results[s.VALUE] = primal_val + data[s.OFFSET]
            new_results[s.PRIMAL] = results_dict['x']
            if not self.is_mip(data):
                new_results[s.EQ_DUAL] = results_dict["y"][0:dims[s.EQ_DIM]]
                new_results[s.INEQ_DUAL] = results_dict["y"][dims[s.EQ_DIM]:]

        return new_results

import numpy,from scipy,import cvxpy,from cvxpy


#! /usr/bin/env python3
# coding: utf-8

import numpy as np
import chainer
import sys, os
from chainer import cuda, serializers
from chainer import functions as chf
from progressbar import progressbar
import librosa
import soundfile as sf
import pickle as pic

from FCA import FCA
from configure import *

class MNMF_DP(FCA):
    """ Blind Speech Enhancement Using Multichannel Nonnegative Matrix Factorization with a Deep Speech Prior (MNMF-DP)

    X_FTM: the observed complex spectrogram
    covarianceMatrix_NFMM: spatial covariance matrices (SCMs) for each source
    W_noise_NnFK: basis vectors for noise sources (Nn means the number of noise sources)
    H_noise_NnKT: activations for noise sources
    Z_speech_DT: latent variables for speech
    power_speech_FT: power spectra of speech that is the output of DNN(Z_speech_DT)
    lambda_NFT: power spectral densities of each source
        lambda_NFT[0] = U_F * V_T * power_speech_FT
        lambda_NFT[1:] = W_noise_NnFK @ H_noise_NnKT
    """

    def __init__(self, speech_VAE=None, n_noise=1, n_Z_iteration=30, n_latent=16, n_basis_noise=2, xp=np, init_SCM="unit", mode_update_parameter=["all", "Z", "one_by_one"][1],\
            mode_update_Z=["sampling", "backprop"][0], normalize_encoder_input=True, seed=0):
        """ initialize MNMF_DP

        Parameters:
        -----------
            speech_VAE: VAE
                trained speech VAE network
            n_noise: int
                the number of noise sources
            n_Z_iteration: int
                the number of iteration for updating Z per global iteration
            n_latent: int
                the dimension of latent variable Z
            n_basis_noise: int
                the number of bases of each noise source
            xp : numpy or cupy
            init_SCM: str
                how to initialize covariance matrix {unit, obs, ILRMA}
            mode_update_parameter: str
                'all' : update all the variables simultanesouly
                'one_by_one' : update one by one
            mode_update_Z: str
                how to update latent variable Z {sampling, backprop}
            normalize_encoder_input: boolean
                normalize observation to initialize latent variable by feeding the observation into a encoder
        """
        super(MNMF_DP, self).__init__(n_source=n_noise+1, xp=xp, init_SCM=init_SCM, mode_update_parameter=mode_update_parameter, seed=seed)
        self.method_name = "MNMF_DP"
        self.n_source, self.n_noise, self.n_speech = n_noise+1, n_noise, 1
        self.n_basis_noise = n_basis_noise
        self.n_Z_iteration = n_Z_iteration
        self.n_latent = n_latent
        self.speech_VAE = speech_VAE
        self.mode_update_Z = mode_update_Z
        self.normalize_encoder_input = normalize_encoder_input

    def initialize_PSD(self):
        self.W_noise_NnFK = self.xp.random.rand(self.n_noise, self.n_freq, self.n_basis_noise).astype(self.xp.float)
        self.H_noise_NnKT = self.xp.random.rand(self.n_noise, self.n_basis_noise, self.n_time).astype(self.xp.float)

        self.u_F = self.xp.ones(self.n_freq, dtype=self.xp.float) / self.n_freq
        self.v_T = self.xp.ones(self.n_time, dtype=self.xp.float)

        power_observation_FT = (self.xp.abs(self.X_FTM) ** 2).mean(axis=2)
        if self.normalize_encoder_input:
            power_observation_FT = power_observation_FT / power_observation_FT.sum(axis=0).mean()
        self.Z_speech_DT = self.speech_VAE.encode_cupy(power_observation_FT.astype(self.xp.float32))
        self.z_link_speech = Z_link(self.Z_speech_DT.T)
        self.z_optimizer_speech = chainer.optimizers.Adam().setup(self.z_link_speech)
        self.power_speech_FT = self.speech_VAE.decode_cupy(self.Z_speech_DT)

        self.lambda_NFT = self.xp.zeros([self.n_source, self.n_freq, self.n_time]).astype(self.xp.float)
        self.lambda_NFT[0] = self.u_F[:, None] * self.v_T[None] * self.power_speech_FT
        self.lambda_NFT[1:] = self.W_noise_NnFK @ self.H_noise_NnKT

    def make_filename_suffix(self):
        self.filename_suffix = f"N={self.n_noise}-it={self.n_iteration}-itZ={self.n_Z_iteration}-Kn={self.n_basis_noise}-D={self.n_latent}-init={self.init_SCM}-latent={self.mode_update_Z}-update={self.mode_update_parameter}"

        if hasattr(self, "name_DNN"):
            self.filename_suffix += f"-DNN={self.name_DNN}"
        if hasattr(self, "file_id"):
            self.filename_suffix += f"-ID={self.file_id}"
        print("param:", self.filename_suffix)

    def update(self):
        if self.mode_update_parameter == "one_by_one":
            self.update_axiliary_variable()
            self.update_W_noise()
            self.update_axiliary_variable()
            self.update_H_noise()
            self.update_axiliary_variable()
            self.update_covarianceMatrix()
            self.update_axiliary_variable()
            self.update_U()
            self.update_axiliary_variable()
            self.update_V()
            self.update_axiliary_variable()
            self.update_Z_speech(calc_constant=True)
            self.normalize()
        elif self.mode_update_parameter == "all":
            self.update_axiliary_variable_and_Z()
            self.update_WH_noise()
            self.update_covarianceMatrix()
            self.update_UV()
            self.update_Z_speech(calc_constant=False)
            self.normalize()
        elif self.mode_update_parameter == "Z":
            self.update_axiliary_variable_and_Z()
            self.update_WH_noise()
            self.update_covarianceMatrix()
            self.update_UV()
            self.update_Z_speech(calc_constant=True)
            self.normalize()

    def update_axiliary_variable_and_Z(self):
        Y_NFTMM = self.lambda_NFT[..., None, None] * self.covarianceMatrix_NFMM[:, :, None]
        if self.xp == np:
            self.Yinv_FTMM = np.linalg.inv(Y_NFTMM.sum(axis=0))
            Yx_FTM1 = self.Yinv_FTMM @ self.X_FTM[..., None]
            self.Yinv_X_Yinv_FTMM = Yx_FTM1 @ Yx_FTM1.conj().transpose(0, 1, 3, 2) # for reducing computational cost in case of CPU
            cov_inv_FMM = np.linalg.inv(self.covarianceMatrix_NFMM[0])
        else:
            self.Yinv_FTMM = self.xp.linalg.inv(Y_NFTMM.sum(axis=0))
            Yx_FTM1 = self.Yinv_FTMM @ self.X_FTM[..., None]
            self.Yinv_X_Yinv_FTMM = Yx_FTM1 @ Yx_FTM1.conj().transpose(0, 1, 3, 2) # for reducing computational cost in case of CPU
            cov_inv_FMM = self.xp.linalg.inv(self.covarianceMatrix_NFMM[0])

        self.tr_Cov_Yinv_X_Yinv_NFT = self.xp.trace(self.covarianceMatrix_NFMM[:, :, None] @ self.Yinv_X_Yinv_FTMM[None], axis1=3, axis2=4).real
        self.tr_Cov_Yinv_NFT = self.xp.trace(self.covarianceMatrix_NFMM[:, :, None] @ self.Yinv_FTMM[None], axis1=3, axis2=4).real

        Phi_FTMM = Y_NFTMM[0] @ self.Yinv_FTMM
        self.tr_Omega_Cov_FT = self.tr_Cov_Yinv_NFT[0]
        self.tr_Cov_Phi_X_Phi_FT = self.xp.trace(cov_inv_FMM[:, None] @ Phi_FTMM @ self.XX_FTMM @ Phi_FTMM.transpose(0, 1, 3, 2).conj(), axis1=2, axis2=3).real
        self.UV_FT = self.u_F[:, None] * self.v_T[None]

    def update_UV(self):
        a_1 = (self.u_F[:, None] * self.power_speech_FT * self.tr_Cov_Yinv_X_Yinv_NFT[0]).sum(axis=0)
        b_1 = (self.u_F[:, None] * self.power_speech_FT * self.tr_Cov_Yinv_NFT[0]).sum(axis=0)

        a_2 = (self.v_T[None] * self.power_speech_FT * self.tr_Cov_Yinv_X_Yinv_NFT[0]).sum(axis=1)
        b_2 = (self.v_T[None] * self.power_speech_FT * self.tr_Cov_Yinv_NFT[0]).sum(axis=1)

        self.v_T *= self.xp.sqrt(a_1 / b_1)
        self.u_F *= self.xp.sqrt(a_2 / b_2)

    def update_U(self):
        a_1 = (self.v_T[None] * self.power_speech_FT * self.tr_Cov_Yinv_X_Yinv_NFT[0]).sum(axis=1)
        b_1 = (self.v_T[None] * self.power_speech_FT * self.tr_Cov_Yinv_NFT[0]).sum(axis=1)
        self.u_F *= self.xp.sqrt(a_1 / b_1)
        self.lambda_NFT[0] = self.u_F[:, None] * self.v_T[None] * self.power_speech_FT

    def update_V(self):
        a_1 = (self.u_F[:, None] * self.power_speech_FT * self.tr_Cov_Yinv_X_Yinv_NFT[0]).sum(axis=0)
        b_1 = (self.u_F[:, None] * self.power_speech_FT * self.tr_Cov_Yinv_NFT[0]).sum(axis=0)
        self.v_T *= self.xp.sqrt(a_1 / b_1)
        self.UV_FT = self.u_F[:, None] * self.v_T[None]
        self.lambda_NFT[0] = self.u_F[:, None] * self.v_T[None] * self.power_speech_FT

    def update_WH_noise(self):
        a_1 = (self.H_noise_NnKT.transpose(0, 2, 1)[:, None] * self.tr_Cov_Yinv_X_Yinv_NFT[1:, :, :, None]).sum(axis=2) # Nn F K
        b_1 = (self.H_noise_NnKT.transpose(0, 2, 1)[:, None] * self.tr_Cov_Yinv_NFT[1:, :, :, None]).sum(axis=2) # Nn F K

        a_2 = (self.W_noise_NnFK[..., None] * self.tr_Cov_Yinv_X_Yinv_NFT[1:, :, None]).sum(axis=1) # Nn K T
        b_2 = (self.W_noise_NnFK[..., None] * self.tr_Cov_Yinv_NFT[1:, :, None]).sum(axis=1) # Nn K T

        self.W_noise_NnFK *= self.xp.sqrt(a_1 / b_1)
        self.H_noise_NnKT *= self.xp.sqrt(a_2 / b_2)

    def update_H_noise(self):
        a_1 = (self.W_noise_NnFK[..., None] * self.tr_Cov_Yinv_X_Yinv_NFT[1:, :, None]).sum(axis=1) # Nn K T
        b_1 = (self.W_noise_NnFK[..., None] * self.tr_Cov_Yinv_NFT[1:, :, None]).sum(axis=1) # Nn K T
        self.H_noise_NnKT *= self.xp.sqrt(a_1 / b_1)
        self.lambda_NFT[1:] = self.W_noise_NnFK @ self.H_noise_NnKT + EPS

    def update_W_noise(self):
        a_1 = (self.H_noise_NnKT.transpose(0, 2, 1)[:, None] * self.tr_Cov_Yinv_X_Yinv_NFT[1:, :, :, None]).sum(axis=2) # Nn F K
        b_1 = (self.H_noise_NnKT.transpose(0, 2, 1)[:, None] * self.tr_Cov_Yinv_NFT[1:, :, :, None]).sum(axis=2) # Nn F K
        self.W_noise_NnFK *= self.xp.sqrt(a_1 / b_1)
        self.lambda_NFT[1:] = self.W_noise_NnFK @ self.H_noise_NnKT + EPS

    def normalize(self):
        mu_NF = self.xp.trace(self.covarianceMatrix_NFMM, axis1=2, axis2=3).real
        self.covarianceMatrix_NFMM = self.covarianceMatrix_NFMM / mu_NF[:, :, None, None]
        self.u_F *= mu_NF[0]
        self.W_noise_NnFK *= mu_NF[1:][:, :, None]

        nu = self.u_F.sum()
        self.u_F /= nu
        self.v_T *= nu

        nu_NnK = self.W_noise_NnFK.sum(axis=1)
        self.W_noise_NnFK /= nu_NnK[:, None]
        self.H_noise_NnKT *= nu_NnK[:, :, None]

        self.lambda_NFT[0] = self.u_F[:, None] * self.v_T[None] * self.power_speech_FT
        self.lambda_NFT[1:] = self.W_noise_NnFK @ self.H_noise_NnKT + EPS

    def loss_func_Z(self, z, vae, n):
        power_FT = chf.exp(vae.decode(z).T) * self.UV_FT + EPS
        if n == 0:
            loss = chf.sum(1 / power_FT * self.tr_Cov_Phi_X_Phi_FT + power_FT * self.tr_Omega_Cov_FT)
        else:
            raise NotImplementedError
        return loss

    def update_Z_speech(self, var_propose_distribution=1e-4, calc_constant=True):
        """
        Parameters:
            var_propose_distribution: float
                the variance of the propose distribution

        Results:
            self.Z_speech_DT: self.xp.array [ n_latent x T ]
                the latent variable of each speech
        """
        if calc_constant:
            self.calculate_constant_for_update_Z()

        if "backprop" in self.mode_update_Z: # acceptance rate is calculated from likelihood
            for it in range(self.n_Z_iteration):
                with chainer.using_config('train', False):
                    self.z_optimizer_speech.update(self.loss_func_Z, self.z_link_speech.z, self.speech_VAE, 0)

            self.Z_speech_DT = self.z_link_speech.z.data.T
            self.power_speech_FT = self.speech_VAE.decode_cupy(self.Z_speech_DT)

        if "sampling" in self.mode_update_Z:
            log_var = self.xp.log(self.xp.ones_like(self.Z_speech_DT).astype(self.xp.float32) * var_propose_distribution)
            Z_speech_old_DT = self.Z_speech_DT
            lambda_speech_old_FT = self.speech_VAE.decode_cupy(Z_speech_old_DT) * self.UV_FT
            for it in range(self.n_Z_iteration):
                Z_speech_new_DT = chf.gaussian(Z_speech_old_DT, log_var).data
                lambda_speech_new_FT = self.speech_VAE.decode_cupy(Z_speech_new_DT) * self.UV_FT
                acceptance_rate =  self.xp.exp((-1 * (1/lambda_speech_new_FT - 1/lambda_speech_old_FT) * self.tr_Cov_Phi_X_Phi_FT -  (lambda_speech_new_FT - lambda_speech_old_FT) * self.tr_Omega_Cov_FT).sum(axis=0) - (Z_speech_new_DT ** 2 - Z_speech_old_DT ** 2).sum(axis=0)/2)
                acceptance_boolean = self.xp.random.random([self.n_time]) < acceptance_rate
                Z_speech_old_DT[:, acceptance_boolean] = Z_speech_new_DT[:, acceptance_boolean]
                lambda_speech_old_FT[:, acceptance_boolean] = lambda_speech_new_FT[:, acceptance_boolean]

            self.Z_speech_DT = Z_speech_old_DT
            self.z_link_speech.z = chainer.Parameter(self.Z_speech_DT.T)
            self.power_speech_FT = self.speech_VAE.decode_cupy(self.Z_speech_DT)

    def calculate_constant_for_update_Z(self):
        Y_NFTMM = self.lambda_NFT[..., None, None] * self.covarianceMatrix_NFMM[:, :, None]
        if self.xp == np:
            self.Yinv_FTMM = np.linalg.inv(Y_NFTMM.sum(axis=0))
            cov_inv_FMM = np.linalg.inv(self.covarianceMatrix_NFMM[0])
        else:
            self.Yinv_FTMM = self.xp.linalg.inv(Y_NFTMM.sum(axis=0))
            cov_inv_FMM = self.xp.linalg.inv(self.covarianceMatrix_NFMM[0])

        Phi_FTMM = Y_NFTMM[0] @ self.Yinv_FTMM
        self.tr_Omega_Cov_FT = self.xp.trace(self.covarianceMatrix_NFMM[0, :, None] @ self.Yinv_FTMM, axis1=2, axis2=3).real
        self.tr_Cov_Phi_X_Phi_FT = self.xp.trace(cov_inv_FMM[:, None] @ Phi_FTMM @ self.XX_FTMM @ Phi_FTMM.transpose(0, 1, 3, 2).conj(), axis1=2, axis2=3).real

    def save_parameter(self, filename):
        param_list = [self.covarianceMatrix_NFMM, self.lambda_NFT, self.u_F, self.v_T, self.Z_speech_DT, self.W_noise_NnFK, self.H_noise_NnKT]

        if self.xp != np:
            param_list = [cuda.to_cpu(param) for param in param_list]

        pic.dump(param_list, open(filename, "wb"))

    def load_parameter(self, filename):
        param_list = pic.load(open(filename, "rb"))
        if self.xp != np:
            param_list = [cuda.to_gpu(param) for param in param_list]

        self.covarianceMatrix_NFMM, self.lambda_NFT, self.u_F, self.v_T, self.Z_speech_DT, self.W_noise_NnFK, self.H_noise_NnKT = param_list
        self.n_source, self.n_freq, self.n_time = self.lambda_NFT.shape
        self.n_mic = self.covarianceMatrix_NFMM.shape[-1]
        self.n_latent = self.Z_speech_DT.shape[0]
        self.n_noise, self.n_speech = self.n_source - 1, 1

class Z_link(chainer.link.Link):
    def __init__(self, z):
        super(Z_link, self).__init__()

        with self.init_scope():
            self.z = chainer.Parameter(z)

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument(         'input_fileName', type= str, help='filename of the multichannel observed signals')
    parser.add_argument(              '--file_id', type= str, default="None", help='file id')
    parser.add_argument(                  '--gpu', type= int, default=     0, help='GPU ID')##
    parser.add_argument(                '--n_fft', type= int, default=  1024, help='number of frequencies')
    parser.add_argument(              '--n_noise', type= int, default=     1, help='number of noise')
    parser.add_argument(             '--n_latent', type= int, default=    16, help='dimention of encoded vector')
    parser.add_argument(                '--n_mic', type= int, default=     8, help='number of microphones')
    parser.add_argument(        '--n_basis_noise', type= int, default=    64, help='number of basis of noise (MODE_noise=NMF)')
    parser.add_argument(             '--init_SCM', type=  str, default="obs", help='unit, obs, ILRMA')
    parser.add_argument(          '--n_iteration', type= int, default=    30, help='number of iteration')
    parser.add_argument(        '--n_Z_iteration', type= int, default=    30, help='number of update Z iteration')
    parser.add_argument(        '--mode_update_Z', type= str, default="sampling", help='sampling, sampling2, backprop, backprop2, hybrid, hybrid2')
    parser.add_argument('--mode_update_parameter', type= str, default= "all", help='all, one_by_one')
    args = parser.parse_args()

    sys.path.append("../DeepSpeechPrior")
    import network_VAE
    model_fileName = f"../DeepSpeechPrior/model-VAE-best-scale=gamma-D={args.n_latent}.npz"
    speech_VAE = network_VAE.VAE(n_latent=args.n_latent)
    serializers.load_npz(model_fileName, speech_VAE)
    name_DNN = "VAE"

    if args.gpu < 0:
        import numpy as xp
    else:
        import cupy as xp
        print("Use GPU " + str(args.gpu))
        cuda.get_device_from_id(args.gpu).use()
        speech_VAE.to_gpu()

    wav, fs = sf.read(args.input_fileName)
    wav = wav.T
    M = min(args.n_mic, len(wav))
    for m in range(M):
        tmp = librosa.core.stft(wav[m], n_fft=args.n_fft, hop_length=int(args.n_fft/4))
        if m == 0:
            spec = np.zeros([tmp.shape[0], tmp.shape[1], M], dtype=np.complex)
        spec[:, :, m] = tmp

    separater = MNMF_DP(n_noise=args.n_noise, n_Z_iteration=args.n_Z_iteration, speech_VAE=speech_VAE, n_latent=args.n_latent, n_basis_noise=args.n_basis_noise, xp=xp, init_SCM=args.init_SCM, mode_update_parameter=args.mode_update_parameter, seed=0)

    separater.load_spectrogram(spec)
    separater.name_DNN = name_DNN
    separater.file_id = args.file_id
    separater.solve(n_iteration=args.n_iteration, save_likelihood=False, save_parameter=False, save_path="./", interval_save_parameter=100)