from __future__ import print_function
import numpy as np
from discretize.utils import mkvc

########################################################################################
#                  DIRECT CURRENT RESISTIVITY AND INDUCED POLARIZATION
########################################################################################

def read_dcip3d_ubc(file_name, data_type):
    """
    Read UBC DCIP3D formatted survey, predicted or observed data files.

    This function loads 3D DC or IP data formatted for the UBC DCIP3D
    coding package and outputs a SimPEG data object. The function will
    recognize if the file is a survey, predicted data or observed data file.

    Parameters
    ----------
    file_name : str
        The file path to the data file
    data_type: str {'volt', 'apparent_chargeability', 'secondary_potential'}
        Defining the input data type:

        - 'volt': DC resistivity data as voltages
        - 'apparent_chargeability': IP data as apparent chargeabilities
        - 'secondary_potential': IP data as secondary potentials

    Returns
    -------
    data
        A SimPEG.data.Data object containing:

        - The survey
        - Observed/predicted data (if present in the data file)
        - Uncertainties (if present in the data file). Note that predicted DC data
        files contain the apparent resistivities, which are loaded into SimPEG and
        defined as uncertainties.

    """
    assert data_type.lower() in [
        "volt",
        "apparent_chargeability",
        "secondary_potential",
    ], "Parameter 'data_type' must be one of {'volt', 'apparent_chargeability', 'secondary_potential'}"

    return _read_dcip_3d_or_octree_ubc(file_name, data_type, "dcip3d")

def read_dcipoctree_ubc(file_name, data_type):
    """
    Read UBC DCIPoctree formatted survey, predicted or observed data files.

    This function loads 3D DC or IP data formatted for the UBC DCIPoctree
    coding package and outputs a SimPEG data object. The function requires
    the user to define whether the data are DC resistivity or IP.

    Parameters
    ----------
    file_name : str
        The file path to the data file
    data_type: str {'volt', 'apparent_chargeability'}
        Defining the input data type:

        - 'volt': DC resistivity data as voltages
        - 'apparent_chargeability': IP data as apparent chargeabilities

    Returns
    -------
    data_object
        A SimPEG.data.Data object containing:

        - The survey
        - Observed/predicted data (if present in the data file)
        - Uncertainties (if present in the data file). Note that predicted DC data
        files contain the apparent resistivities, which are loaded into SimPEG and
        defined as uncertainties.

    """

    # Unused for now but it will be when we manage IP types better.
    assert data_type.lower() in [
        "volt",
        "apparent_chargeability",
    ], "Parameter 'data_type' must be one of {'volt', 'apparent_chargeability'}"

    return _read_dcip_3d_or_octree_ubc(file_name, data_type, "dcipoctree")

def _read_dcip_3d_or_octree_ubc(file_name, data_type, code_type):
    """
    Read 3D DC/IP survey, predicted and observation files in UBC-GIF format.

    Parameters
    ----------

    file_name : str
        The file path to the data file
    code_type : str {'dcip3d', 'dcipoctree'}
        Code type. Choose from {'dcip3d', 'dcipoctree'}

    Returns
    -------
    data_object
        A SimPEG.data.Data object containing:

        - The survey
        - Observed/predicted data (if present in the data file)
        - Uncertainties (if present in the data file). Note that predicted DC data
        files contain the apparent resistivities, which are loaded into SimPEG and
        defined as uncertainties.

    """
    assert data_type.lower() in [
        "volt",
        "apparent_chargeability",
        "secondary_potential",
    ], "Parameter 'data_type' must be one of {'volt', 'apparent_chargeability', 'secondary_potential'}"

    assert code_type.lower() in [
        "dcip3d",
        "dcipoctree",
    ], "Parameter 'code_type' must be one of {'dcip3d', 'dcipoctree'}"

    # Prevent circular import
    from ...electromagnetics.static import resistivity as dc
    from ...data import Data

    # Load file
    obsfile = np.genfromtxt(file_name, delimiter=" \n", dtype=np.str, comments="!")

    # Pre-allocate
    source_list = []
    receiver_list = []
    d = []
    wd = []

    # Flag for z value provided
    is_surface = False
    is_pole_tx = False
    is_pole_rx = False

    # IP data for dcip3d has a line with a flag we can remove.
    if (code_type == "dcip3d") & (data_type != "volt"):
        obsfile = obsfile[1:]

    # Since SimPEG defines secondary potential from IP as voltage,
    # we must use this type when defining the receivers.
    if data_type.lower() == "secondary_potential":
        data_type = "volt"

    # Countdown for number of obs/tx
    count = 0
    for ii in range(obsfile.shape[0]):

        if not obsfile[ii]:
            continue

        # Extract transmitter ii and the number of receivers
        if count == 0:
            rx = []
            temp = np.fromstring(obsfile[ii], dtype=float, sep=" ").T
            count = int(temp[-1])

            # Check if z value is provided, if False -> nan
            if len(temp) == 5:
                # check if pole|dipole
                if np.allclose(temp[0:2], temp[2:4]):
                    tx = np.r_[temp[0:2], np.nan]
                    is_pole_tx = True
                else:
                    tx = np.r_[temp[0:2], np.nan, temp[2:4], np.nan]
                is_surface = True

            else:
                # check if pole|dipole
                if np.allclose(temp[0:3], temp[3:6]):
                    tx = np.r_[temp[0:3]]
                    is_pole_tx = True
                else:
                    tx = temp[:-1]

            continue

        # Extract receivers
        temp = np.fromstring(obsfile[ii], dtype=float, sep=" ")

        if is_surface:
            data_column_index = 4  # Since dpred for dc has app_res

            # Check if Pole Receiver
            if np.allclose(temp[0:2], temp[2:4]):
                is_pole_rx = True
                rx.append(temp[:2])
            else:
                rx.append(np.r_[temp[0:2], np.nan, temp[2:4], np.nan])

        else:
            data_column_index = 6  # Since dpred for dc has app_res

            # Check if Pole Receiver
            if np.allclose(temp[0:3], temp[3:6]):
                is_pole_rx = True
                rx.append(temp[:3])
            else:
                rx.append(temp[:6])

        # Predicted/observed data
        if len(temp) == data_column_index + 1:
            d.append(temp[data_column_index])

        # Observed data or predicted DC data (since app res column)
        elif len(temp) == data_column_index + 2:
            d.append(temp[data_column_index])
            wd.append(temp[data_column_index + 1])

        count = count - 1

        # Reach the end of transmitter block
        if count == 0:
            rx = np.asarray(rx)
            if is_pole_rx:
                Rx = dc.receivers.Pole(rx[:, :3], data_type=data_type)
            else:
                Rx = dc.receivers.Dipole(rx[:, :3], rx[:, 3:], data_type=data_type)
            if is_pole_tx:
                source_list.append(dc.sources.Pole([Rx], tx[:3]))
            else:
                source_list.append(dc.sources.Dipole([Rx], tx[:3], tx[3:]))

    # Define survey type
    if is_pole_tx:
        str1 = "pole-"
    else:
        str1 = "dipole-"

    if is_pole_rx:
        str2 = "pole"
    else:
        str2 = "dipole"

    electrode_configuration = str1 + str2
    survey = dc.survey.Survey(source_list, survey_type=electrode_configuration)
    data_out = Data(survey=survey)

    if len(d) > 0:
        data_out.dobs = d

    if len(wd) > 0:
        data_out.standard_deviation = wd

    return data_out

def write_dcip3d_ubc(
    file_name,
    data_object,
    data_type,
    file_type,
    format_type="general",
    electrode_configuration=None,
    comment_lines="",
):
    """
    Write UBC DCIP3D formatted survey, predicted or observation files.

    Parameters
    ----------
    file_name:
    data_object:
    file_type: 'survey', 'dpred', 'dobs'
    format_type: 'general', 'surface', 'simple'
    data_type: 'volt', 'apparent_chargeability', 'secondary_potential'
    electrode_configuration: 'pole-pole', 'pole-dipole', 'dipole-pole', 'dipole-dipole'
    comment_lines:)
    """

    assert data_type.lower() in [
        "volt",
        "apparent_chargeability",
        "secondary_potential",
    ], "Parameter 'data_type' must be one of {'volt', 'apparent_chargeability', 'secondary_potential'}"

    assert file_type.lower() in [
        "survey",
        "dpred",
        "dobs",
    ], "Parameter 'file_type' must be one of {'survey', 'dpred', 'dobs'}"

    _write_dcip_3d_or_octree_ubc(
        file_name,
        data_object,
        data_type,
        file_type,
        format_type=format_type,
        electrode_configuration=electrode_configuration,
        code_type="dcip3d",
        comment_lines=comment_lines,
    )

def write_dcipoctree_ubc(
    file_name,
    data_object,
    data_type,
    file_type,
    format_type="general",
    electrode_configuration=None,
    comment_lines="",
):
    """
    Write UBC DCIPoctree formatted survey, predicted or observation files.

    Parameters
    ----------
    file_name:
    data_object:
    file_type: 'survey', 'dpred', 'dobs'
    format_type: 'general', 'surface', 'simple'
    data_type: 0 (DC), 1 (IP), 2 (another IP)
    electrode_configuration: 'pole-pole', 'pole-dipole', 'dipole-pole', 'dipole-dipole'
    comment_lines:)
    """

    assert data_type.lower() in [
        "volt",
        "apparent_chargeability",
    ], "Parameter 'data_type' must be one of {'volt', 'apparent_chargeability'}"

    assert file_type.lower() in [
        "survey",
        "dpred",
        "dobs",
    ], "Parameter 'file_type' must be one of {'survey', 'dpred', 'dobs'}"

    _write_dcip_3d_or_octree_ubc(
        file_name,
        data_object,
        data_type,
        file_type,
        format_type="general",
        electrode_configuration=electrode_configuration,
        code_type="dcipoctree",
        comment_lines=comment_lines,
    )

def _write_dcip_3d_or_octree_ubc(
    file_name,
    data_object,
    data_type,
    file_type,
    format_type="general",
    electrode_configuration=None,
    code_type="dcip3d",
    comment_lines=None,
):
    """
    Write UBC DCIP3D formatted survey, predicted or observation files.

    Parameters
    ----------
    file_name:
    data_object:
    file_type: 'survey', 'dpred', 'dobs'
    format_type: 'general', 'surface', 'simple'
    data_type: {'volt', 'apparent_chargeability', 'secondary_potential'}
    electrode_configuration: 'pole-pole', 'pole-dipole', 'dipole-pole', 'dipole-dipole'
    code_type: 'dcip3d', 'dcipoctree'
    comment_lines:
    """

    # Prevent circular import
    from ...electromagnetics.static import resistivity as dc
    from ...electromagnetics.static.utils.static_utils import apparent_resistivity
    from ...data import Data

    # Validate inputs
    if not isinstance(data_object, Data):
        raise Exception(
            "A Data instance ({datacls}: <{datapref}.{datacls}>) must be "
            "provided as the second input. The provided input is a "
            "{providedcls} <{providedpref}.{providedcls}>".format(
                datacls=Data.__name__,
                datapref=Data.__module__,
                providedcls=data_object.__class__.__name__,
                providedpref=data_object.__module__,
            )
        )

    assert data_type.lower() in [
        "volt",
        "apparent_chargeability",
        "secondary_potential",
    ], "Parameter 'data_type' must be one of {'volt', 'apparent_chargeability', 'secondary_potential'}"

    assert file_type.lower() in [
        "survey",
        "dpred",
        "dobs",
    ], "Parameter 'file_type' must be one of {'survey', 'dpred', 'dobs'}"

    assert code_type.lower() in [
        "dcip3d",
        "dcipoctree",
    ], "Parameter 'code_type' must be one of {'dcip3d', 'dcipoctree'}"

    format_type = format_type.lower()
    if format_type not in ["surface", "general", "simple"]:
        raise Exception(
            "format_type must be 'surface' | 'general' | 'simple' "
            " not {}".format(format_type)
        )

    if electrode_configuration is None:
        electrode_configuration = data_object.survey.survey_type

    # Predicted DC data will automatically contain apparent resistivity column.
    # Here we compute the apparent resistivities and treat it like an uncertainties column.
    if (file_type.lower() == "dpred") & (data_type == "volt"):
        data_object.standard_deviation = apparent_resistivity(data_object)
        file_type = "dobs"

    # Write comments and IP type (if applicable)
    with open(file_name, "w") as fid:
        fid.write(f"! {format_type} FORMAT\n")

        if comment_lines is not None and len(comment_lines) > 0:
            # ensure comment_lines ends with a new line character
            if comment_lines[-1] != "\n":
                comment_lines += "\n"
            fid.write(comment_lines)

        # DCIP3D will allow user to choose definition of IP data. DC data has no flag.
        # DCIPoctree IP data is always apparent chargeability.
        if (code_type.lower() == "dcip3d") & (data_type == "apparent_chargeability"):
            fid.write("IPTYPE=1\n")
        elif (code_type.lower() == "dcip3d") & (data_type == "secondary_potential"):
            fid.write("IPTYPE=2\n")

        # Index deciding if z locations are written
        if format_type.lower() == "surface":
            end_index = 2
        elif format_type.lower() == "general":
            end_index = 3

        # Loop over all sources
        count = 0
        for src in data_object.survey.source_list:

            # Write Source
            nD = src.nD

            if electrode_configuration.lower() in ["pole-dipole", "pole-pole"]:
                tx = np.r_[src.location]
                tx = np.repeat(np.r_[[tx]], 2, axis=0)
            elif electrode_configuration.lower() in ["dipole-dipole", "dipole-pole"]:
                tx = np.c_[src.location]

            fid.writelines("%e " % ii for ii in mkvc(tx[:, 0:end_index].T))
            fid.write(f"{nD}\n")

            # Write receivers
            for rx in src.receiver_list:

                if electrode_configuration.lower() in ["pole-dipole", "dipole-dipole"]:
                    M = rx.locations[0][0:end_index]
                    N = rx.locations[1][0:end_index]
                elif electrode_configuration.lower() in ["pole-pole", "dipole-pole"]:
                    M = rx.locations[0:end_index]
                    N = rx.locations[0:end_index]

                if file_type.lower() != "survey":
                    N = np.c_[N, data_object.dobs[count : count + rx.nD]]

                if file_type.lower() == "dobs":
                    N = np.c_[N, data_object.standard_deviation[count : count + rx.nD]]

                # Write receivers and locations
                if isinstance(N, np.ndarray):
                    np.savetxt(
                        fid, np.c_[M, N], fmt="%e",
                    )
                else:
                    raise Exception(
                        """Uncertainities SurveyObject.std should be set.
                        Either float or nunmpy.ndarray is expected, """
                        "not {}".format(type(data_object.relative_error))
                    )

                fid.write("\n")

                count += rx.nD


import numpy
import conf
import sys
import random
def STATS(NEVENT,NGOOD):
	# IMPLICIT #real*8 (A-H,O-Z)
	# IMPLICIT #integer*8 (I-N)
	# COMMON/SETP/TMAX,SMALL,API,ESTART,THETA,PHI,TCFMAX(10),TCFMAX1,RSTART,EFIELD,ETHRM,ECUT,NDELTA,IMIP,IWRITE
	# COMMON/IDEXC/NGEXC1,NGEXC2,NGEXC3,NGEXC4,NGEXC5,NGEXC6,IDG1,IDG2,IDG3,IDG4,IDG5,IDG6
	# COMMON/PRIM3/MSUM(10000),MCOMP(10000),MRAYL(10000),MPAIR(10000),MPHOT(10000),MVAC(10000)
	# COMMON/STTS/XST(150000),YST(150000),ZST(150000),TIME(150000),TTIME(150000),NFGF(150000),NFGPP(150000),NFGBR(150000),NELEC,NEGION,EST1,EST2
	# COMMON/STEXC/XSTEXC(150000),YSTEXC(150000),ZSTEXC(150000),TSTEXC(150000),NSTEXC
	# COMMON/STEXCNUL/XSTN(150000),YSTN(150000),ZSTN(150000),TSTN(150000),IDNUL(150000),NEXCNUL
	# COMMON/CLUS/XAV[100000],YAV[100000],ZAV[100000],TAV[100000],XYAV[100000],XYZAV[100000],DX(100000),DY(100000),DZ[100000],DT(100000),DXY(100000),DXYZ[100000],NCL(100000),FARX1(100000),FARY1(100000),FARZ1(100000),FARXY1(100000),RMAX1(100000),TSUM(100000),XNEG(100000), YNEG(100000),ZNEG(100000),EDELTA[100000],EDELTA2(100000),NCLEXC(100000)
	# COMMON/PLOT/NXPL10(31),NYPL10(31),NZPL10(31),NXPL40(31),NYPL40(31),NZPL40(31),NXPL100(31),NYPL100(31),NZPL100(31),NXPL400(31),NYPL400(31),NZPL400(31),NXPL1000(31),NYPL1000(31),NZPL1000(31),NXPL2(31),NYPL2(31),NZPL2(31),NXPL4000(31),NYPL4000(31),NZPL4000(31),NXPL10000(31),NYPL10000(31),NZPL10000(31),NXPL40000(31),NYPL40000(31),NZPL40000(31),NXPL100000(31),NYPL100000(31),NZPL100000(31),NRPL2(31),NRPL10(31),NRPL40(31),NRPL100(31),NRPL400(31),NRPL1000(31),NRPL4000(31),NRPL10000(31),NRPL40000(31),NRPL100000(31),NEPL1(100),NEPL10(100),NEPL100(100),MELEC(1000),MELEC3(1000),MELEC10(1000),MELEC30(1000),MELEC100(1000),MELEC300(1000)
	#COMMON/SETP/
	global TMAX,SMALL,API,ESTART,THETA,PHI,TCFMAX#(10)
	global TCFMAX1,RSTART,EFIELD,ETHRM,ECUT,NDELTA,IMIP,IWRITE
	#COMMON/IDEXC/
	global NGEXC1,NGEXC2,NGEXC3,NGEXC4,NGEXC5,NGEXC6,IDG1,IDG2,IDG3,IDG4,IDG5,IDG6
	#COMMON/PRIM3/
	global MSUM#(10000)
	global MCOMP#(10000)
	global MRAYL#(10000)
	global MPAIR#(10000)
	global MPHOT#(10000)
	global MVAC#(10000)
	#COMMON/STTS/
	global XST#(150000)
	global YST#(150000)
	global ZST#(150000)
	global TIME#(150000)
	global TTIME#(150000)
	global NFGF#(150000)
	global NFGPP#(150000)
	global NFGBR#(150000)
	global NELEC,NEGION,EST1,EST2
	#COMMON/STEXC/
	global XSTEXC#(150000)
	global YSTEXC#(150000)
	global ZSTEXC#(150000)
	global TSTEXC#(150000)
	global NSTEXC
	#COMMON/STEXCNUL/
	global XSTN#(150000)
	global YSTN#(150000)
	global ZSTN#(150000)
	global TSTN#(150000)
	global IDNUL#(150000)
	global NEXCNUL
	#COMMON/CLUS/
	global XAV#(100000)
	global YAV#(100000)
	global ZAV#(100000)
	global TAV#(100000)
	global XYAV#(100000)
	global XYZAV#(100000)
	global DX#(100000)
	global DY#(100000)
	global DZ#(100000)
	global DT#(100000)
	global DXY#(100000)
	global DXYZ#(100000)
	global NCL#(100000)
	global FARX1#(100000)
	global FARY1#(100000)
	global FARZ1#(100000)
	global FARXY1#(100000)
	global RMAX1#(100000)
	global TSUM#(100000)
	global XNEG#(100000)
	global YNEG#(100000)
	global ZNEG#(100000)
	global EDELTA#(100000)
	global EDELTA2#(100000)
	global NCLEXC#(100000)
	#COMMON/PLOT/
	global NXPL10#(31)
	global NYPL10#(31)
	global NZPL10#(31)
	global NXPL40#(31)
	global NYPL40#(31)
	global NZPL40#(31)
	global NXPL100#(31)
	global NYPL100#(31)
	global NZPL100#(31)
	global NXPL400#(31)
	global NYPL400#(31)
	global NZPL400#(31)
	global NXPL1000#(31)
	global NYPL1000#(31)
	global NZPL1000#(31)
	global NXPL2#(31)
	global NYPL2#(31)
	global NZPL2#(31)
	global NXPL4000#(31)
	global NYPL4000#(31)
	global NZPL4000#(31)
	global NXPL10000#(31)
	global NYPL10000#(31)
	global NZPL10000#(31)
	global NXPL40000#(31)
	global NYPL40000#(31)
	global NZPL40000#(31)
	global NXPL100000#(31)
	global NYPL100000#(31)
	global NZPL100000#(31)
	global NRPL2#(31)
	global NRPL10#(31)
	global NRPL40#(31)
	global NRPL100#(31)
	global NRPL400#(31)
	global NRPL1000#(31)
	global NRPL4000#(31)
	global NRPL10000#(31)
	global NRPL40000#(31)
	global NRPL100000#(31)
	global NEPL1#(100)
	global NEPL10#(100)
	global NEPL100#(100)
	global MELEC#(1000)
	global MELEC3#(1000)
	global MELEC10#(1000)
	global MELEC30#(1000)
	global MELEC100#(1000)
	global MELEC300#(1000)
	def get_globals():
		TMAX=conf.TMAX
		SMALL=conf.SMALL
		API=conf.API
		ESTART=conf.ESTART
		THETA=conf.THETA
		PHI=conf.PHI
		TCFMAX=conf.TCFMAX
		TCFMAX1=conf.TCFMAX1
		RSTART=conf.RSTART
		EFIELD=conf.EFIELD
		ETHRM=conf.ETHRM
		ECUT=conf.ECUT
		NDELTA=conf.NDELTA
		IMIP=conf.IMIP
		IWRITE  =conf.IWRITE
		NGEXC1=conf.NGEXC1
		NGEXC2=conf.NGEXC2
		NGEXC3=conf.NGEXC3
		NGEXC4=conf.NGEXC4
		NGEXC5=conf.NGEXC5
		NGEXC6=conf.NGEXC6
		IDG1=conf.IDG1
		IDG2=conf.IDG2
		IDG3=conf.IDG3
		IDG4=conf.IDG4
		IDG5=conf.IDG5
		IDG6=conf.IDG6
		MSUM=conf.MSUM
		MCOMP=conf.MCOMP
		MRAYL=conf.MRAYL
		MPAIR=conf.MPAIR
		MPHOT=conf.MPHOT
		MVAC=conf.MVAC
		XST=conf.XST
		YST=conf.YST
		ZST=conf.ZST
		TIME=conf.TIME
		TTIME=conf.TTIME
		NFGF=conf.NFGF
		NFGPP=conf.NFGPP
		NFGBR=conf.NFGBR
		NELEC=conf.NELEC
		NEGION=conf.NEGION
		EST1=conf.EST1
		EST2=conf.EST2
		XSTEXC=conf.XSTEXC
		YSTEXC=conf.YSTEXC
		ZSTEXC=conf.ZSTEXC
		TSTEXC=conf.TSTEXC
		NSTEXC=conf.NSTEXC
		XSTN=conf.XSTN
		YSTN=conf.YSTN
		ZSTN=conf.ZSTN
		TSTN=conf.TSTN
		IDNUL=conf.IDNUL
		NEXCNUL=conf.NEXCNUL
		XAV=conf.XAV
		YAV=conf.YAV
		ZAV=conf.ZAV
		TAV=conf.TAV
		XYAV=conf.XYAV
		XYZAV=conf.XYZAV
		DX=conf.DX
		DY=conf.DY
		DZ=conf.DZ
		DT=conf.DT
		DXY=conf.DXY
		DXYZ=conf.DXYZ
		NCL=conf.NCL
		FARX1=conf.FARX1
		FARY1=conf.FARY1
		FARZ1=conf.FARZ1
		FARXY1=conf.FARXY1
		RMAX1=conf.RMAX1
		TSUM=conf.TSUM
		XNEG=conf.XNEG
		YNEG=conf.YNEG
		ZNEG=conf.ZNEG
		EDELTA=conf.EDELTA
		EDELTA2=conf.EDELTA2
		NCLEXC=conf.NCLEXC
		NXPL10=conf.NXPL10
		NYPL10=conf.NYPL10
		NZPL10=conf.NZPL10
		NXPL40=conf.NXPL40
		NYPL40=conf.NYPL40
		NZPL40=conf.NZPL40
		NXPL100=conf.NXPL100
		NYPL100=conf.NYPL100
		NZPL100=conf.NZPL100
		NXPL400=conf.NXPL400
		NYPL400=conf.NYPL400
		NZPL400=conf.NZPL400
		NXPL1000=conf.NXPL1000
		NYPL1000=conf.NYPL1000
		NZPL1000=conf.NZPL1000
		NXPL2=conf.NXPL2
		NYPL2=conf.NYPL2
		NZPL2=conf.NZPL2
		NXPL4000=conf.NXPL4000
		NYPL4000=conf.NYPL4000
		NZPL4000=conf.NZPL4000
		NXPL10000=conf.NXPL10000
		NYPL10000=conf.NYPL10000
		NZPL10000=conf.NZPL10000
		NXPL40000=conf.NXPL40000
		NYPL40000=conf.NYPL40000
		NZPL40000=conf.NZPL40000
		NXPL100000=conf.NXPL100000
		NYPL100000=conf.NYPL100000
		NZPL100000=conf.NZPL100000
		NRPL2=conf.NRPL2
		NRPL10=conf.NRPL10
		NRPL40=conf.NRPL40
		NRPL100=conf.NRPL100
		NRPL400=conf.NRPL400
		NRPL1000=conf.NRPL1000
		NRPL4000=conf.NRPL4000
		NRPL10000=conf.NRPL10000
		NRPL40000=conf.NRPL40000
		NRPL100000=conf.NRPL100000
		NEPL1=conf.NEPL1
		NEPL10=conf.NEPL10
		NEPL100=conf.NEPL100
		MELEC=conf.MELEC
		MELEC3=conf.MELEC3
		MELEC10=conf.MELEC10
		MELEC30=conf.MELEC30
		MELEC100=conf.MELEC100
		MELEC300=conf.MELEC300
		globals().update(locals())
	get_globals()
	def update_globals():
		conf.TMAX=TMAX
		conf.SMALL=SMALL
		conf.API=API
		conf.ESTART=ESTART
		conf.THETA=THETA
		conf.PHI=PHI
		conf.TCFMAX=TCFMAX
		conf.TCFMAX1=TCFMAX1
		conf.RSTART=RSTART
		conf.EFIELD=EFIELD
		conf.ETHRM=ETHRM
		conf.ECUT=ECUT
		conf.NDELTA=NDELTA
		conf.IMIP=IMIP
		conf.IWRITE  =IWRITE
		conf.NGEXC1=NGEXC1
		conf.NGEXC2=NGEXC2
		conf.NGEXC3=NGEXC3
		conf.NGEXC4=NGEXC4
		conf.NGEXC5=NGEXC5
		conf.NGEXC6=NGEXC6
		conf.IDG1=IDG1
		conf.IDG2=IDG2
		conf.IDG3=IDG3
		conf.IDG4=IDG4
		conf.IDG5=IDG5
		conf.IDG6=IDG6
		conf.MSUM=MSUM
		conf.MCOMP=MCOMP
		conf.MRAYL=MRAYL
		conf.MPAIR=MPAIR
		conf.MPHOT=MPHOT
		conf.MVAC=MVAC
		conf.XST=XST
		conf.YST=YST
		conf.ZST=ZST
		conf.TIME=TIME
		conf.TTIME=TTIME
		conf.NFGF=NFGF
		conf.NFGPP=NFGPP
		conf.NFGBR=NFGBR
		conf.NELEC=NELEC
		conf.NEGION=NEGION
		conf.EST1=EST1
		conf.EST2=EST2
		conf.XSTEXC=XSTEXC
		conf.YSTEXC=YSTEXC
		conf.ZSTEXC=ZSTEXC
		conf.TSTEXC=TSTEXC
		conf.NSTEXC=NSTEXC
		conf.XSTN=XSTN
		conf.YSTN=YSTN
		conf.ZSTN=ZSTN
		conf.TSTN=TSTN
		conf.IDNUL=IDNUL
		conf.NEXCNUL=NEXCNUL
		conf.XAV=XAV
		conf.YAV=YAV
		conf.ZAV=ZAV
		conf.TAV=TAV
		conf.XYAV=XYAV
		conf.XYZAV=XYZAV
		conf.DX=DX
		conf.DY=DY
		conf.DZ=DZ
		conf.DT=DT
		conf.DXY=DXY
		conf.DXYZ=DXYZ
		conf.NCL=NCL
		conf.FARX1=FARX1
		conf.FARY1=FARY1
		conf.FARZ1=FARZ1
		conf.FARXY1=FARXY1
		conf.RMAX1=RMAX1
		conf.TSUM=TSUM
		conf.XNEG=XNEG
		conf.YNEG=YNEG
		conf.ZNEG=ZNEG
		conf.EDELTA=EDELTA
		conf.EDELTA2=EDELTA2
		conf.NCLEXC=NCLEXC
		conf.NXPL10=NXPL10
		conf.NYPL10=NYPL10
		conf.NZPL10=NZPL10
		conf.NXPL40=NXPL40
		conf.NYPL40=NYPL40
		conf.NZPL40=NZPL40
		conf.NXPL100=NXPL100
		conf.NYPL100=NYPL100
		conf.NZPL100=NZPL100
		conf.NXPL400=NXPL400
		conf.NYPL400=NYPL400
		conf.NZPL400=NZPL400
		conf.NXPL1000=NXPL1000
		conf.NYPL1000=NYPL1000
		conf.NZPL1000=NZPL1000
		conf.NXPL2=NXPL2
		conf.NYPL2=NYPL2
		conf.NZPL2=NZPL2
		conf.NXPL4000=NXPL4000
		conf.NYPL4000=NYPL4000
		conf.NZPL4000=NZPL4000
		conf.NXPL10000=NXPL10000
		conf.NYPL10000=NYPL10000
		conf.NZPL10000=NZPL10000
		conf.NXPL40000=NXPL40000
		conf.NYPL40000=NYPL40000
		conf.NZPL40000=NZPL40000
		conf.NXPL100000=NXPL100000
		conf.NYPL100000=NYPL100000
		conf.NZPL100000=NZPL100000
		conf.NRPL2=NRPL2
		conf.NRPL10=NRPL10
		conf.NRPL40=NRPL40
		conf.NRPL100=NRPL100
		conf.NRPL400=NRPL400
		conf.NRPL1000=NRPL1000
		conf.NRPL4000=NRPL4000
		conf.NRPL10000=NRPL10000
		conf.NRPL40000=NRPL40000
		conf.NRPL100000=NRPL100000
		conf.NEPL1=NEPL1
		conf.NEPL10=NEPL10
		conf.NEPL100=NEPL100
		conf.MELEC=MELEC
		conf.MELEC3=MELEC3
		conf.MELEC10=MELEC10
		conf.MELEC30=MELEC30
		conf.MELEC100=MELEC100
		conf.MELEC300=MELEC300
	#-----------------------------------------------------------------------
	#  NEVENT IS EVENT EVENT POINTER. NGOOD IS GOOD EVENT POINTER
	#  FORMS AVERAGES OVER EACH DELTA AND DOES SOME STATISTICS
	#  LOADS PLOT ARRAYS XPLOT YPLOT AND ZPLOT (SCALED BY 1 10 AND 100)
	#  OUTPUTS RAW DATA TO FILE IF IWRITE CONTROL GT  0
	#  OUTPUTS THERMALISED ELECTRON X,Y,Z AND T IF  IWRITE EQ 1
	#  OUTPUTS ALSO X,Y,Z AND T FOR EACH  EXCITATION IF
	#  IWRITE EQ 2
	# ----------------------------------------------------------------------
	if(NGOOD <= 0):
		print(' IN STATS NEVENT=',NEVENT,' NGOOD=',NGOOD)
	# endif
	NCLUS=NELEC-NEGION
	if(NCLUS > 150000):
		# GO TO 99
		pass
	else:
		SUMX=0.00
		SUMX2=0.00
		SUMY=0.00
		SUMY2=0.00
		SUMZ=0.00
		SUMZ2=0.00
		SUMRXY=0.00
		SUMRXY2=0.00
		SUMRXYZ=0.00
		SUMRXYZ2=0.00
		SUMT=0.00
		SUMT2=0.00
		FARX=0.00
		FARY=0.00
		FARZ=0.00
		FARXY=0.00
		RMAX=0.00
		SUMTT=0.00
		NXNEG=0
		NYNEG=0
		NZNEG=0
		ESUM=0.0
		ATOTR=0.0
		ATOTC=0.0
		ATOTP=0.0
		ATOTPE=0.0
		#
		for IS in range(1,int(NCLUS)+1):
			XST[IS]=XST[IS]*1.e6
			X=XST[IS]
			if(X < 0.0):
				NXNEG=NXNEG+1
				I1=int(X/2.0-0.5)
				I2=int(X/10.0-0.5)
				I3=int(X/40.0-0.5)
				I4=int(X/100.0-0.5)
				I5=int(X/400.0-0.5)
				I6=int(X/1000.0-0.5)
				I7=int(X/4000.0-0.5)
				I8=int(X/10000.0-0.5)
				I9=int(X/40000.0-0.5)
				I10=int(X/100000.0-0.5)
			else:
				I1=int(X/2.0+0.5)
				I2=int(X/10.0+0.5)
				I3=int(X/40.0+0.5)
				I4=int(X/100.0+0.5)
				I5=int(X/400.0+0.5)
				I6=int(X/1000.0+0.5)
				I7=int(X/4000.0+0.5)
				I8=int(X/10000.0+0.5)
				I9=int(X/40000.0+0.5)
				I10=int(X/100000.0+0.5)
			# endif
			I1=I1+16
			I2=I2+16
			I3=I3+16
			I4=I4+16
			I5=I5+16
			I6=I6+16
			I7=I7+16
			I8=I8+16
			I9=I9+16
			I10=I10+16
			if(I1 < 1):
				I1=1
			if(I1 > 31):
				I1=31
			if(I2 < 1):
				I2=1
			if(I2 > 31):
				I2=31
			if(I3 < 1):
				I3=1
			if(I3 > 31):
				I3=31
			if(I4 < 1):
				I4=1
			if(I4 > 31):
				I4=31
			if(I5 < 1):
				I5=1
			if(I5 > 31):
				I5=31
			if(I6 < 1):
				I6=1
			if(I6 > 31):
				I6=31
			if(I7 < 1):
				I7=1
			if(I7 > 31):
				I7=31
			if(I8 < 1):
				I8=1
			if(I8 > 31):
				I8=31
			if(I9 < 1):
				I9=1
			if(I9 > 31):
				I9=31
			if(I10 < 1):
				I10=1
			if(I10 > 31):
				I10=31
			NXPL2[I1]=NXPL2[I1]+1
			NXPL10[I2]=NXPL10[I2]+1
			NXPL40[I3]=NXPL40[I3]+1
			NXPL100[I4]=NXPL100[I4]+1
			NXPL400[I5]=NXPL400[I5]+1
			NXPL1000[I6]=NXPL1000[I6]+1
			NXPL4000[I7]=NXPL4000[I7]+1
			NXPL10000[I8]=NXPL10000[I8]+1
			NXPL40000[I9]=NXPL40000[I9]+1
			NXPL100000[I10]=NXPL100000[I10]+1
			X2=X*X
			SUMX=SUMX+X
			SUMX2=SUMX2+X2
			if(abs(X)> abs(FARX)):
				FARX=abs(X)
			YST[IS]=YST[IS]*1.e6
			Y=YST[IS]
			if(Y < 0.0):
				NYNEG=NYNEG+1
				I1=int(Y/2.0-0.5)
				I2=int(Y/10.0-0.5)
				I3=int(Y/40.0-0.5)
				I4=int(Y/100.0-0.5)
				I5=int(Y/400.0-0.5)
				I6=int(Y/1000.0-0.5)
				I7=int(Y/4000.0-0.5)
				I8=int(Y/10000.0-0.5)
				I9=int(Y/40000.0-0.5)
				I10=int(Y/100000.0-0.5)
			else:
				I1=int(Y/2.0+0.5)
				I2=int(Y/10.0+0.5)
				I3=int(Y/40.0+0.5)
				I4=int(Y/100.0+0.5)
				I5=int(Y/400.0+0.5)
				I6=int(Y/1000.0+0.5)
				I7=int(Y/4000.0+0.5)
				I8=int(Y/10000.0+0.5)
				I9=int(Y/40000.0+0.5)
				I10=int(Y/100000.0+0.5)
			# endif
			I1=I1+16
			I2=I2+16
			I3=I3+16
			I4=I4+16
			I5=I5+16
			I6=I6+16
			I7=I7+16
			I8=I8+16
			I9=I9+16
			I10=I10+16
			if(I1 < 1):
				I1=1
			if(I1 > 31):
				I1=31
			if(I2 < 1):
				I2=1
			if(I2 > 31):
				I2=31
			if(I3 < 1):
				I3=1
			if(I3 > 31):
				I3=31
			if(I4 < 1):
				I4=1
			if(I4 > 31):
				I4=31
			if(I5 < 1):
				I5=1
			if(I5 > 31):
				I5=31
			if(I6 < 1):
				I6=1
			if(I6 > 31):
				I6=31
			if(I7 < 1):
				I7=1
			if(I7 > 31):
				I7=31
			if(I8 < 1):
				I8=1
			if(I8 > 31):
				I8=31
			if(I9 < 1):
				I9=1
			if(I9 > 31):
				I9=31
			if(I10 < 1):
				I10=1
			if(I10 > 31):
				I10=31
			NYPL2[I1]=NYPL2[I1]+1
			NYPL10[I2]=NYPL10[I2]+1
			NYPL40[I3]=NYPL40[I3]+1
			NYPL100[I4]=NYPL100[I4]+1
			NYPL400[I5]=NYPL400[I5]+1
			NYPL1000[I6]=NYPL1000[I6]+1
			NYPL4000[I7]=NYPL4000[I7]+1
			NYPL10000[I8]=NYPL10000[I8]+1
			NYPL40000[I9]=NYPL40000[I9]+1
			NYPL100000[I10]=NYPL100000[I10]+1
			Y2=Y*Y
			SUMY=SUMY+Y
			SUMY2=SUMY2+Y2
			if(abs(Y)> abs(FARY)):
				FARY=abs(Y)
			ZST[IS]=ZST[IS]*1.e6
			Z=ZST[IS]
			if(Z < 0.0):
				NZNEG=NZNEG+1
				I1=int(Z/2.0-0.5)
				I2=int(Z/10.0-0.5)
				I3=int(Z/40.0-0.5)
				I4=int(Z/100.0-0.5)
				I5=int(Z/400.0-0.5)
				I6=int(Z/1000.0-0.5)
				I7=int(Z/4000.0-0.5)
				I8=int(Z/10000.0-0.5)
				I9=int(Z/40000.0-0.5)
				I10=int(Z/100000.0-0.5)
			else:
				I1=int(Z/2.0+0.5)
				I2=int(Z/10.0+0.5)
				I3=int(Z/40.0+0.5)
				I4=int(Z/100.0+0.5)
				I5=int(Z/400.0+0.5)
				I6=int(Z/1000.0+0.5)
				I7=int(Z/4000.0+0.5)
				I8=int(Z/10000.0+0.5)
				I9=int(Z/40000.0+0.5)
				I10=int(Z/100000.0+0.5)
			# endif
			I1=I1+16
			I2=I2+16
			I3=I3+16
			I4=I4+16
			I5=I5+16
			I6=I6+16
			I7=I7+16
			I8=I8+16
			I9=I9+16
			I10=I10+16
			if(I1 < 1):
				I1=1
			if(I1 > 31):
				I1=31
			if(I2 < 1):
				I2=1
			if(I2 > 31):
				I2=31
			if(I3 < 1):
				I3=1
			if(I3 > 31):
				I3=31
			if(I4 < 1):
				I4=1
			if(I4 > 31):
				I4=31
			if(I5 < 1):
				I5=1
			if(I5 > 31):
				I5=31
			if(I6 < 1):
				I6=1
			if(I6 > 31):
				I6=31
			if(I7 < 1):
				I7=1
			if(I7 > 31):
				I7=31
			if(I8 < 1):
				I8=1
			if(I8 > 31):
				I8=31
			if(I9 < 1):
				I9=1
			if(I9 > 31):
				I9=31
			if(I10 < 1):
				I10=1
			if(I10 > 31):
				I10=31
			NZPL2[I1]=NZPL2[I1]+1
			NZPL10[I2]=NZPL10[I2]+1
			NZPL40[I3]=NZPL40[I3]+1
			NZPL100[I4]=NZPL100[I4]+1
			NZPL400[I5]=NZPL400[I5]+1
			NZPL1000[I6]=NZPL1000[I6]+1
			NZPL4000[I7]=NZPL4000[I7]+1
			NZPL10000[I8]=NZPL10000[I8]+1
			NZPL40000[I9]=NZPL40000[I9]+1
			NZPL100000[I10]=NZPL100000[I10]+1
			Z=ZST[IS]
			R=math.sqrt(X*X+Y*Y+Z*Z)
			I1=int(R/2.0+0.5)
			I2=int(R/10.0+0.5)
			I3=int(R/40.0+0.5)
			I4=int(R/100.0+0.5)
			I5=int(R/400.0+0.5)
			I6=int(R/1000.0+0.5)
			I7=int(R/4000.0+0.5)
			I8=int(R/10000.0+0.5)
			I9=int(R/40000.0+0.5)
			I10=int(R/100000.0+0.5)
			I1=I1+16
			I2=I2+16
			I3=I3+16
			I4=I4+16
			I5=I5+16
			I6=I6+16
			I7=I7+16
			I8=I8+16
			I9=I9+16
			I10=I10+16
			if(I1 > 31):
				I1=31
			if(I2 > 31):
				I2=31
			if(I3 > 31):
				I3=31
			if(I4 > 31):
				I4=31
			if(I5 > 31):
				I5=31
			if(I6 > 31):
				I6=31
			if(I7 > 31):
				I7=31
			if(I8 > 31):
				I8=31
			if(I9 > 31):
				I9=31
			if(I10 > 31):
				I10=31
			if(I1 < 1):
				I1=1
			if(I2 < 1):
				I2=1
			if(I3 < 1):
				I3=1
			if(I4 < 1):
				I4=1
			if(I5 < 1):
				I5=1
			if(I6 < 1):
				I6=1
			if(I7 < 1):
				I7=1
			if(I8 < 1):
				I8=1
			if(I9 < 1):
				I9=1
			if(I10 < 1):
				I10=1
			NRPL2[I1]=NRPL2[I1]+1
			NRPL10[I2]=NRPL10[I2]+1
			NRPL40[I3]=NRPL40[I3]+1
			NRPL100[I4]=NRPL100[I4]+1
			NRPL400[I5]=NRPL400[I5]+1
			NRPL1000[I6]=NRPL1000[I6]+1
			NRPL4000[I7]=NRPL4000[I7]+1
			NRPL10000[I8]=NRPL10000[I8]+1
			NRPL40000[I9]=NRPL40000[I9]+1
			NRPL100000[I10]=NRPL100000[I10]+1
			Z2=Z*Z
			SUMZ=SUMZ+Z
			SUMZ2=SUMZ2+Z2
			if(abs(Z)> abs(FARZ)):
				FARZ=abs(Z)
			RXY=math.sqrt(X2+Y2)
			RXYZ=math.sqrt(X2+Y2+Z2)
			SUMRXY=SUMRXY+RXY
			SUMRXY2=SUMRXY2+RXY*2
			SUMRXYZ=SUMRXYZ+RXYZ
			SUMRXYZ2=SUMRXYZ2+RXYZ*2
			if(RXY > FARXY):
				FARXY=RXY
			if(RXYZ > RMAX):
				RMAX=RXYZ
			T=TIME[IS]
			SUMT=SUMT+T
			SUMT2=SUMT2+T*T
			SUMTT=SUMTT+TTIME(IS)
			#
			XSTEXC[IS]=XSTEXC[IS]*1.e6
			YSTEXC[IS]=YSTEXC[IS]*1.e6
			ZSTEXC[IS]=ZSTEXC[IS]*1.e6
		# 400 CONTINUE
		# OUTPUT THERMAL ELECTRON POSITIONS AND TIME
		f=open("fort.50","w+")
		if(IWRITE == 1):
			#      ITEST=123456
			#      WRITE(50,*) ITEST
			# WRITE(50,*)
			f.write(NGOOD,NCLUS,NSTEXC,NGEXC1,NGEXC2,NGEXC3,NGEXC4,NGEXC5,NGEXC6,MCOMP[NEVENT],MPAIR[NEVENT],NEXCNUL)
			#      DO 8876 IPR=1,NCLUS
			#      WRITE(50,*) IPR,XST[IPR],YST[IPR],ZST[IPR],TIME[IPR],NFGF[IPR],
			#    /NFGPP[IPR],NFGBR[IPR]
			#8876 CONTINUE
			# WRITE(50,*)
			for IPR in range(1,NCLUS+1):

				f.write(XST[IPR],YST[IPR],ZST[IPR],TIME[IPR],NFGF[IPR],NFGPP[IPR],NFGBR[IPR])
		# endif
		# OUTPUT EXCITATION CLOUD COORDINATES HERE
		if(IWRITE == 2):
			#      ITEST=123456
			#      WRITE(50,*) ITEST
			# WRITE(50,*)
			f.write(NGOOD,NCLUS,NSTEXC,NGEXC1,NGEXC2,NGEXC3,NGEXC4,NGEXC5,NGEXC6,MCOMP[NEVENT],MPAIR[NEVENT],NEXCNUL)
			#      DO 8877 IPR=1,NCLUS
			#      WRITE(50,*) IPR,XST[IPR],YST[IPR],ZST[IPR],TIME[IPR],NFGF[IPR],
			#    /NFGPP[IPR],NFGBR[IPR]
			#8877 CONTINUE
			# WRITE(50,*)
			for IPR in range(1,NCLUS+1):
				f.write(XST[IPR],YST[IPR],ZST[IPR],TIME[IPR],NFGF[IPR],NFGPP[IPR],NFGBR[IPR])
			#       DO 8878 IPR=1,NSTEXC
			#       WRITE(50,*) IPR,XSTEXC[IPR],YSTEXC[IPR],ZSTEXC[IPR],TSTEXC[IPR]
			# 8878 CONTINUE
			# WRITE(50,*)
			for IPR in range(1,NSTEXC+1):
				f.write(XSTEXC[IPR],YSTEXC[IPR],ZSTEXC[IPR],TSTEXC[IPR])
			if(NEXCNUL > 0):
				# WRITE(50,*)
				for IPR in range(1,NEXCNUL+1):
					f.write(XSTN[IPR],YSTN[IPR],ZSTN[IPR],TSTN[IPR],IDNUL[IPR])
			else:
				# WRITE BLANK LINE
				# WRITE(50,*)
				f.write(0.0,0.0,0.0,0.0,0 )
			# endif
		# endif
		f.close()
		#---------------------------------------------------
		I1=int(EST1+1.0)
		I2=int(EST1/10.0+1.0)
		I3=int(EST1/100.0+1.0)
		if(I1 > 100):
			I1=100
		if(I2 > 100):
			I2=100
		if(I3 > 100):
			I3=100
		NEPL1[I1]=NEPL1[I1]+1
		NEPL10[I2]=NEPL10[I2]+1
		NEPL100[I3]=NEPL100[I3]+1
		KDUM=NELEC
		KDUM3=1+(NELEC/3)
		KDUM10=1+(NELEC/10)
		KDUM30=1+(NELEC/30)
		KDUM100=1+(NELEC/100)
		KDUM300=1+(NELEC/300)
		if(KDUM > 1000):
			KDUM=1000
		MELEC[int(KDUM)]=MELEC[int(KDUM)]+1
		if(KDUM3 > 1000):
			KDUM3=1000
		MELEC3[int(KDUM3)]=MELEC3[int(KDUM3)]+1
		if(KDUM10 > 1000):
			KDUM10=1000
		MELEC10[int(KDUM10)]=MELEC10[int(KDUM10)]+1
		if(KDUM30 > 1000):
			KDUM30=1000
		MELEC30[int(KDUM30)]=MELEC30[int(KDUM30)]+1
		if(KDUM100 > 1000):
			KDUM100=1000
		MELEC100[int(KDUM100)]=MELEC100[int(KDUM100)]+1
		if(KDUM300 > 1000):
			KDUM300=1000
		MELEC300[int(KDUM300)]=MELEC300[int(KDUM300)]+1
		#
		# STORE AVERAGES AND WIDTHS FOR EACH DELTA
		#
		if(NCLUS == 0):
			return
		ACLUS=float(NCLUS)
		XAV[NGOOD]=SUMX/ACLUS
		YAV[NGOOD]=SUMY/ACLUS
		ZAV[NGOOD]=SUMZ/ACLUS
		TAV[NGOOD]=SUMT/ACLUS
		XYAV[NGOOD]=SUMRXY/ACLUS
		XYZAV[NGOOD]=SUMRXYZ/ACLUS
		#  IONISATION CLUSTER SIZE
		NCL[NGOOD]=NCLUS
		# EXCITATION CLUSTER SIZE
		NCLEXC[NGOOD]=NSTEXC
		#
		FARX1[NGOOD]=FARX
		FARY1[NGOOD]=FARY
		FARZ1[NGOOD]=FARZ
		FARXY1[NGOOD]=FARXY
		RMAX1[NGOOD]=RMAX
		TSUM[NGOOD]=SUMTT
		XNEG[NGOOD]=float(NXNEG)/ACLUS
		YNEG[NGOOD]=float(NYNEG)/ACLUS
		ZNEG[NGOOD]=float(NZNEG)/ACLUS
		EDELTA[NGOOD]=EST1
		EDELTA2[NGOOD]=EST2
		if(NCLUS > 1):
			NC2=NCLUS*NCLUS-NCLUS
			ANC2=float(NC2)
			DX[NGOOD]=math.sqrt(abs((ACLUS*SUMX2-SUMX*SUMX)/ANC2))
			DY[NGOOD]=math.sqrt(abs((ACLUS*SUMY2-SUMY*SUMY)/ANC2))
			DZ[NGOOD]=math.sqrt(abs((ACLUS*SUMZ2-SUMZ*SUMZ)/ANC2))
			DT[NGOOD]=math.sqrt(abs((ACLUS*SUMT2-SUMT*SUMT)/ANC2))
			DXY[NGOOD]=math.sqrt(abs((ACLUS*SUMRXY2-SUMRXY*SUMRXY)/ANC2))
			DXYZ[NGOOD]=math.sqrt(abs((ACLUS*SUMRXYZ2-SUMRXYZ*SUMRXYZ)/ANC2))
		else:
			DX[NGOOD]=0.0
			DY[NGOOD]=0.0
			DZ[NGOOD]=0.0
			DZ[NGOOD]=0.0
			DT[NGOOD]=0.0
			DXY[NGOOD]=0.0
			DXYZ[NGOOD]=0.0
		# endif
		return
	# 99 WRITE(6,991) NCLUS
	# 991
	print('\n\n\n WARNING OVERFLOW IN ARRAYS IN FUNCTION STATS. NCLUS =',NCLUS,' STOPPED: FUNCTION')
	sys.exit()
	# end


"""Molecular database (MDB) class

   * MdbExomol is the MDB for ExoMol
   * MdbHit is the MDB for HITRAN or HITEMP

"""
import numpy as np
import jax.numpy as jnp
import pathlib
from exojax.spec import hapi, exomolapi, exomol
from exojax.spec.hitran import gamma_natural as gn
import pandas as pd

__all__ = ['MdbExomol','MdbHit']

class MdbExomol(object):
    """ molecular database of ExoMol

    MdbExomol is a class for ExoMol.

    Attributes:
        nurange: nu range [min,max] (cm-1)
        nu_lines (nd array): line center (cm-1)
        Sij0 (nd array): line strength at T=Tref (cm)
        dev_nu_lines (jnp array): line center in device (cm-1)
        logsij0 (jnp array): log line strength at T=Tref
        A (jnp array): Einstein A coeeficient
        gamma_natural (jnp array): gamma factor of the natural broadening
        elower (jnp array): the lower state energy (cm-1)
        gpp (jnp array): statistical weight
        jlower (jnp array): J_lower
        jupper (jnp array): J_upper
        n_Tref (jnp array): temperature exponent
        alpha_ref (jnp array): alpha_ref (gamma0)
        n_Tref_def: default temperature exponent in .def file, used for jlower not given in .broad
        alpha_ref_def: default alpha_ref (gamma0) in .def file, used for jlower not given in .broad

    """
    def __init__(self,path,nurange=[-np.inf,np.inf],margin=1.0,crit=-np.inf, bkgdatm="H2", broadf=True):
        """Molecular database for Exomol form

        Args:
           path: path for Exomol data directory/tag. For instance, "/home/CO/12C-16O/Li2015"
           nurange: wavenumber range list (cm-1) [min,max] or wavenumber grid
           margin: margin for nurange (cm-1)
           crit: line strength lower limit for extraction
           bkgdatm: background atmosphere for broadening. e.g. H2, He,
           broadf: if False, the default broadening parameters in .def file is used

        Note:
           The trans/states files can be very large. For the first time to read it, we convert it to the feather-format. After the second-time, we use the feather format instead.

        """
        explanation="Note: Couldn't find the feather format. We convert data to the feather format. After the second time, it will become much faster."

        self.path = pathlib.Path(path)
        t0=self.path.parents[0].stem
        molec=t0+"__"+str(self.path.stem)
        self.bkgdatm=bkgdatm
        print("Background atmosphere: ",self.bkgdatm)
        molecbroad=t0+"__"+self.bkgdatm

        self.crit = crit
        self.margin = margin
        self.nurange=[np.min(nurange),np.max(nurange)]
        self.broadf=broadf
        #Where exomol files are
        self.states_file = self.path/pathlib.Path(molec+".states.bz2")
        self.pf_file = self.path/pathlib.Path(molec+".pf")
        self.def_file = self.path/pathlib.Path(molec+".def")
        self.broad_file = self.path/pathlib.Path(molecbroad+".broad")

        if not self.def_file.exists():
                self.download(molec,extension=[".def"])
        if not self.pf_file.exists():
                self.download(molec,extension=[".pf"])
        if not self.states_file.exists():
                self.download(molec,extension=[".states.bz2"])
        if not self.broad_file.exists():
                self.download(molec,extension=[".broad"])

        #load def
        self.n_Texp_def, self.alpha_ref_def, self.molmass, numinf, numtag=exomolapi.read_def(self.def_file)
        #  default n_Texp value if not given
        if self.n_Texp_def is None:
            self.n_Texp_def=0.5
        #  default alpha_ref value if not given
        if self.alpha_ref_def is None:
            self.alpha_ref_def=0.07

        #load states
        if self.states_file.with_suffix(".feather").exists():
            states=pd.read_feather(self.states_file.with_suffix(".feather"))
        else:
            print(explanation)
            states=exomolapi.read_states(self.states_file)
            states.to_feather(self.states_file.with_suffix(".feather"))
        #load pf
        pf=exomolapi.read_pf(self.pf_file)
        self.gQT=jnp.array(pf["QT"].to_numpy()) #grid QT
        self.T_gQT=jnp.array(pf["T"].to_numpy()) #T forgrid QT

        #trans file(s)
        print("Reading transition file")
        if numinf is None:
            self.trans_file = self.path/pathlib.Path(molec+".trans.bz2")
            if not self.trans_file.exists():
                self.download(molec,[".trans.bz2"])

            if self.trans_file.with_suffix(".feather").exists():
                trans=pd.read_feather(self.trans_file.with_suffix(".feather"))
                ndtrans=trans.to_numpy()
                del trans
            else:
                print(explanation)
                trans=exomolapi.read_trans(self.trans_file)
                trans.to_feather(self.trans_file.with_suffix(".feather"))
                ndtrans=trans.to_numpy()
                del trans
            #compute gup and elower
            self._A, self.nu_lines, self._elower, self._gpp, self._jlower, self._jupper=exomolapi.pickup_gE(states,ndtrans)
        else:
            imin=np.searchsorted(numinf,self.nurange[0],side="right")-1 #left side
            imax=np.searchsorted(numinf,self.nurange[1],side="right")-1 #left side
            self.trans_file=[]
            for k,i in enumerate(range(imin,imax+1)):
                trans_file = self.path/pathlib.Path(molec+"__"+numtag[i]+".trans.bz2")
                if not trans_file.exists():
                    self.download(molec,extension=[".trans.bz2"],numtag=numtag[i])
                if trans_file.with_suffix(".feather").exists():
                    trans=pd.read_feather(trans_file.with_suffix(".feather"))
                    ndtrans=trans.to_numpy()
                    del trans
                else:
                    print(explanation)
                    trans=exomolapi.read_trans(trans_file)
                    trans.to_feather(trans_file.with_suffix(".feather"))
                    ndtrans=trans.to_numpy()
                    del trans
                self.trans_file.append(trans_file)
                #compute gup and elower
                if k==0:
                    self._A, self.nu_lines, self._elower, self._gpp, self._jlower, self._jupper=exomolapi.pickup_gE(states,ndtrans)
                else:
                    Ax, nulx, elowerx, gppx, jlowerx, jupperx=exomolapi.pickup_gE(states,ndtrans)
                    self._A=np.hstack([self._A,Ax])
                    self.nu_lines=np.hstack([self.nu_lines,nulx])
                    self._elower=np.hstack([self._elower,elowerx])
                    self._gpp=np.hstack([self._gpp,gppx])
                    self._jlower=np.hstack([self._jlower,jlowerx])
                    self._jupper=np.hstack([self._jupper,jupperx])

        self.Tref=296.0
        self.QTref=np.array(self.QT_interp(self.Tref))

        ##Line strength: input should be ndarray not jnp array
        self.Sij0=exomol.Sij0(self._A,self._gpp,self.nu_lines,self._elower,self.QTref)

        ### MASKING ###
        mask=(self.nu_lines>self.nurange[0]-self.margin)\
        *(self.nu_lines<self.nurange[1]+self.margin)\
        *(self.Sij0>self.crit)

        self.masking(mask)

    def masking(self,mask):
        """applying mask and (re)generate jnp.arrays

        Args:
           mask: mask to be applied. self.mask is updated.

        Note:
           We have nd arrays and jnp arrays. We apply the mask to nd arrays and generate jnp array from the corresponding nd array. For instance, self._A is nd array and self.A is jnp array.

        """
        #numpy float 64 Do not convert them jnp array
        self.nu_lines = self.nu_lines[mask]
        self.Sij0 = self.Sij0[mask]
        self._A=self._A[mask]
        self._elower=self._elower[mask]
        self._gpp=self._gpp[mask]
        self._jlower=self._jlower[mask]
        self._jupper=self._jupper[mask]

        #jnp arrays
        self.dev_nu_lines=jnp.array(self.nu_lines)
        self.logsij0=jnp.array(np.log(self.Sij0))
        self.A=jnp.array(self._A)
        self.gamma_natural=gn(self.A)
        self.elower=jnp.array(self._elower)
        self.gpp=jnp.array(self._gpp)
        self.jlower=jnp.array(self._jlower,dtype=int)
        self.jupper=jnp.array(self._jupper,dtype=int)
        ##Broadening parameters
        self.set_broadening()

    def set_broadening(self,alpha_ref_def=None,n_Texp_def=None):
        """setting broadening parameters

        Args:
           alpha_ref: set default alpha_ref and apply it. None=use self.alpha_ref_def
           n_Texp_def: set default n_Texp and apply it. None=use self.n_Texp_def
        """
        if alpha_ref_def:
            self.alpha_ref_def = alpha_ref_def
        if n_Texp_def:
            self.n_Texp_def = n_Texp_def

        if self.broadf:
            try:
                print(".broad is used.")
                bdat=exomolapi.read_broad(self.broad_file)
                codelv=exomolapi.check_bdat(bdat)
                print("Broadening code level=",codelv)
                if codelv=="a0":
                    j2alpha_ref, j2n_Texp = exomolapi.make_j2b(bdat,\
                        alpha_ref_default=self.alpha_ref_def,\
                        n_Texp_default=self.n_Texp_def,\
                        jlower_max=np.max(self._jlower))
                    self.alpha_ref=jnp.array(j2alpha_ref[self._jlower])
                    self.n_Texp=jnp.array(j2n_Texp[self._jlower])
                elif codelv=="a1":
                    j2alpha_ref, j2n_Texp = exomolapi.make_j2b(bdat,\
                        alpha_ref_default=self.alpha_ref_def,\
                        n_Texp_default=self.n_Texp_def,\
                        jlower_max=np.max(self._jlower))
                    jj2alpha_ref, jj2n_Texp=exomolapi.make_jj2b(bdat,\
                        j2alpha_ref_def=j2alpha_ref,j2n_Texp_def=j2n_Texp,\
                        jupper_max=np.max(self._jupper))
                    self.alpha_ref=jnp.array(jj2alpha_ref[self._jlower,self._jupper])
                    self.n_Texp=jnp.array(jj2n_Texp[self._jlower,self._jupper])
            except:
                print("Warning: Cannot load .broad. The default broadening parameters are used.")
                self.alpha_ref=jnp.array(self.alpha_ref_def*np.ones_like(self._jlower))
                self.n_Texp=jnp.array(self.n_Texp_def*np.ones_like(self._jlower))

        else:
            print("The default broadening parameters are used.")
            self.alpha_ref=jnp.array(self.alpha_ref_def*np.ones_like(self._jlower))
            self.n_Texp=jnp.array(self.n_Texp_def*np.ones_like(self._jlower))

    def QT_interp(self,T):
        """interpolated partition function

        Args:
           T: temperature

        Returns:
           Q(T) interpolated in jnp.array

        """
        return jnp.interp(T,self.T_gQT,self.gQT)

    def qr_interp(self,T):
        """interpolated partition function ratio

        Args:
           T: temperature

        Returns:
           qr(T)=Q(T)/Q(Tref) interpolated in jnp.array

        """
        return self.QT_interp(T)/self.QT_interp(self.Tref)

    def download(self,molec,extension,numtag=None):
        """Downloading Exomol files

        Args:
           molec: like "12C-16O__Li2015"
           extension: extension list e.g. [".pf",".def",".trans.bz2",".states.bz2",".broad"]
           numtag: number tag of transition file if exists. e.g. "11100-11200"

        Note:
           The download URL is written in exojax.utils.url.

        """
        import urllib.request
        from exojax.utils.molname import e2s
        import os
        from exojax.utils.url import url_ExoMol

        tag=molec.split("__")
        molname_simple=e2s(tag[0])

        for ext in extension:
            if ext==".trans.bz2" and numtag is not None:
                ext="__"+numtag+ext

            if ext==".broad":
                pfname_arr=[tag[0]+"__H2"+ext,tag[0]+"__He"+ext,tag[0]+"__air"+ext]
                url = url_ExoMol()+molname_simple+"/"+tag[0]+"/"
            else:
                pfname_arr=[molec+ext]
                url = url_ExoMol()+molname_simple+"/"+tag[0]+"/"+tag[1]+"/"

            for pfname in pfname_arr:
                pfpath=url+pfname
                os.makedirs(str(self.path), exist_ok=True)
                print("Downloading "+pfpath)
                try:
                    urllib.request.urlretrieve(pfpath,str(self.path/pfname))
                except:
                    print("Error: Couldn't download "+ext+" file and save.")

class MdbHit(object):
    """ molecular database of ExoMol

    MdbExomol is a class for ExoMol.

    Attributes:
        nurange: nu range [min,max] (cm-1)
        nu_lines (nd array): line center (cm-1)
        Sij0 (nd array): line strength at T=Tref (cm)
        dev_nu_lines (jnp array): line center in device (cm-1)
        logsij0 (jnp array): log line strength at T=Tref
        A (jnp array): Einstein A coeeficient
        gamma_natural (jnp array): gamma factor of the natural broadening
        gamma_air (jnp array): gamma factor of air pressure broadening
        gamma_self (jnp array): gamma factor of self pressure broadening
        elower (jnp array): the lower state energy (cm-1)
        gpp (jnp array): statistical weight
        n_air (jnp array): air temperature exponent

    """

    def __init__(self,path,nurange=[-np.inf,np.inf],margin=250.0,crit=-np.inf):
        """Molecular database for HITRAN/HITEMP form

        Args:
           path: path for HITRAN/HITEMP par file
           nurange: wavenumber range list (cm-1) [min,max] or wavenumber grid
           margin: margin for nurange (cm-1)
           crit: line strength lower limit for extraction

        """
        #downloading
        self.path = pathlib.Path(path)
        molec=str(self.path.stem)
        if not self.path.exists():
            self.download()

        #bunzip2 if suffix is .bz2
        if self.path.suffix==".bz2":
            import bz2,shutil
            if self.path.with_suffix('').exists():
                import os
                os.remove(self.path.with_suffix(''))
            print("bunziping")
            with bz2.BZ2File(str(self.path)) as fr:
                with open(str(self.path.with_suffix('')),"wb") as fw:
                    shutil.copyfileobj(fr,fw)
            self.path=self.path.with_suffix('')

        molec=str(self.path.stem)

        hapi.db_begin(str(self.path.parent))
        self.Tref=296.0
        self.molecid = search_molecid(molec)
        self.crit = crit
        self.margin = margin
        self.nurange=[np.min(nurange),np.max(nurange)]

        #nd arrays using DRAM (not jnp, not in GPU)
        self.nu_lines = hapi.getColumn(molec, 'nu')
        self.Sij0 = hapi.getColumn(molec, 'sw')
        self.delta_air = hapi.getColumn(molec, 'delta_air')
        self.isoid = hapi.getColumn(molec,'local_iso_id')
        self.uniqiso=np.unique(self.isoid)

        self._A=hapi.getColumn(molec, 'a')
        self._n_air = hapi.getColumn(molec, 'n_air')
        self._gamma_air = hapi.getColumn(molec, 'gamma_air')
        self._gamma_self =hapi.getColumn(molec, 'gamma_self')
        self._elower = hapi.getColumn(molec, 'elower')
        self._gpp = hapi.getColumn(molec, 'gpp')

        ### MASKING ###
        mask=(self.nu_lines>self.nurange[0]-self.margin)\
        *(self.nu_lines<self.nurange[1]+self.margin)\
        *(self.Sij0>self.crit)

        self.masking(mask)

    def masking(self,mask):
        """applying mask and (re)generate jnp.arrays

        Args:
           mask: mask to be applied

        Note:
           We have nd arrays and jnp arrays. We apply the mask to nd arrays and generate jnp array from the corresponding nd array. For instance, self._A is nd array and self.A is jnp array.

        """

        #numpy float 64 Do not convert them jnp array
        self.nu_lines = self.nu_lines[mask]
        self.Sij0 = self.Sij0[mask]
        self.delta_air=self.delta_air[mask]
        self.isoid = self.isoid[mask]
        self.uniqiso=np.unique(self.isoid)

        ##numpy float 64 copy source for jnp
        self._A=self._A[mask]
        self._n_air = self._n_air[mask]
        self._gamma_air = self._gamma_air[mask]
        self._gamma_self = self._gamma_self[mask]
        self._elower = self._elower[mask]
        self._gpp = self._gpp[mask]

        #jnp.array copy from the copy sources
        self.dev_nu_lines=jnp.array(self.nu_lines)
        self.logsij0=jnp.array(np.log(self.Sij0))
        self.A=jnp.array(self._A)
        self.n_air=jnp.array(self._n_air)
        self.gamma_air = jnp.array(self._gamma_air)
        self.gamma_self = jnp.array(self._gamma_self)
        self.elower=jnp.array(self._elower)
        self.gpp=jnp.array(self._gpp)
        self.gamma_natural=gn(self.A)

    def download(self):
        """Downloading HITRAN/HITEMP par file

        Note:
           The download URL is written in exojax.utils.url.

        """
        import urllib.request
        from exojax.utils.url import url_HITRAN12
        from exojax.utils.url import url_HITEMP

        try:
            url = url_HITRAN12()+self.path.name
            urllib.request.urlretrieve(url,str(self.path))
        except:
            print(url)
            print("HITRAN download failed")
        try:
            url = url_HITEMP()+self.path.name
            print(url)
            urllib.request.urlretrieve(url,str(self.path))
        except:
            print("HITEMP download failed")

    ####################################

    def ExomolQT(self,path):
        """use a partition function from ExoMol

        Args:
           path: path for Exomol data directory/tag. For instance, "/home/CO/12C-16O/Li2015"

        """
        #load pf

        self.empath = pathlib.Path(path)
        t0=self.empath.parents[0].stem
        molec=t0+"__"+str(self.empath.stem)
        self.pf_file = self.empath/pathlib.Path(molec+".pf")
        if not self.pf_file.exists():
                self.exomol_pf_download(molec)

        pf=exomolapi.read_pf(self.pf_file)
        self.gQT=jnp.array(pf["QT"].to_numpy()) #grid QT
        self.T_gQT=jnp.array(pf["T"].to_numpy()) #T forgrid QT

    def exomol_pf_download(self,molec):
        """Downloading Exomol pf files

        Args:
           molec: like "12C-16O__Li2015"

        Note:
           The download URL is written in exojax.utils.url.

        """
        import urllib.request
        from exojax.utils.molname import e2s
        import os
        from exojax.utils.url import url_ExoMol

        tag=molec.split("__")
        molname_simple=e2s(tag[0])
        url = url_ExoMol()+molname_simple+"/"+tag[0]+"/"+tag[1]+"/"

        ext=".pf"
        pfname=molec+ext
        pfpath=url+pfname
        os.makedirs(str(self.empath), exist_ok=True)
        print("Downloading "+pfpath)
        try:
            urllib.request.urlretrieve(pfpath,str(self.empath/pfname))
        except:
            print("Error: Couldn't download "+ext+" file and save.")

    def QT_interp(self,T):
        """interpolated partition function

        Args:
           T: temperature

        Returns:
           Q(T) interpolated in jnp.array

        """
        return jnp.interp(T,self.T_gQT,self.gQT)

    def qr_interp(self,T):
        """interpolated partition function ratio

        Args:
           T: temperature

        Returns:
           qr(T)=Q(T)/Q(Tref) interpolated in jnp.array

        """
        return self.QT_interp(T)/self.QT_interp(self.Tref)

    def Qr_HAPI(self,Tarr):
        """Partition Function ratio using HAPI partition sum

        Args:
           Tarr: temperature array (K)

        Returns:
           Qr = partition function ratio array [N_Tarr x N_iso]

        Note:
           N_Tarr = len(Tarr), N_iso = len(self.uniqiso)

        """
        allT=list(np.concatenate([[self.Tref],Tarr]))
        Qrx=[]
        for iso in self.uniqiso:
            Qrx.append(hapi.partitionSum(self.molecid,iso, allT))
        Qrx=np.array(Qrx)
        qr=Qrx[:,1:].T/Qrx[:,0] #Q(T)/Q(Tref)
        return qr

    def Qr_line_HAPI(self,T):
        """Partition Function ratio using HAPI partition sum

        Args:
           T: temperature (K)

        Returns:
           Qr_line, partition function ratio array for lines [Nlines]

        Note:
           Nlines=len(self.nu_lines)

        """
        qr_line=np.ones_like(self.isoid,dtype=np.float64)
        qrx=self.Qr_HAPI([T])
        for idx,iso in enumerate(self.uniqiso):
            mask=self.isoid==iso
            qr_line[mask]=qrx[0,idx]
        return qr_line

    def Qr_layer_HAPI(self,Tarr):
        """Partition Function ratio using HAPI partition sum

        Args:
           Tarr: temperature array (K)

        Returns:
           Qr_layer, partition function ratio array for lines [N_Tarr x Nlines]

        Note:
           Nlines=len(self.nu_lines)
           N_Tarr=len(Tarr)

        """
        NP=len(Tarr)
        qt=np.zeros((NP,len(self.isoid)))
        qr=self.Qr_HAPI(Tarr)
        for idx,iso in enumerate(self.uniqiso):
            mask=self.isoid==iso
            for ilayer in range(NP):
                qt[ilayer,mask]=qr[ilayer,idx]
        return qt

def search_molecid(molec):
    """molec id from molec (source table name) of HITRAN/HITEMP

    Args:
       molec: source table name

    Return:
       int: molecid (HITRAN molecular id)

    """
    try:
        hitf=molec.split("_")
        molecid=int(hitf[0])
        return molecid

    except:
        print("Warning: Define molecid by yourself.")
        return None

if __name__ == "__main__":
    #mdb=MdbExomol("/home/kawahara/exojax/data/CO/12C-16O/Li2015/")
    #mdb=MdbExomol("/home/kawahara/exojax/data/CH4/12C-1H4/YT34to10/",nurange=[6050.0,6150.0])
    mdb=MdbExomol('.database/H2O/1H2-16O/POKAZATEL',[4310.0,4320.0],crit=1.e-45)

#    mask=mdb.A>1.e-42
#    mdb.masking(mask)
#    mdb=MdbExomol("/home/kawahara/exojax/data/exomol/NH3/14N-1H3/CoYuTe/",nurange=[6050.0,6150.0])
#    mdb=MdbExomol("/home/kawahara/exojax/data/exomol/H2S/1H2-32S/AYT2/",nurange=[6050.0,6150.0])
#    mdb=MdbExomol("/home/kawahara/exojax/data/exomol/FeH/56Fe-1H/MoLLIST/",nurange=[6050.0,6150.0])
#    mdb=MdbExomol("/home/kawahara/exojax/data/exomol/NO/14N-16O/NOname/14N-16O__NOname")


# Copyright 2021 D-Wave Systems Inc.
#
#    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.

r"""Contains implementations for applying MPO's to MPS's with subsequent
compression.
"""

#####################################
## Load libraries/packages/modules ##
#####################################

# For deep copies of objects.
import copy

# For general array handling.
import numpy as np

# For calculating the eigenspectra of Hermitian matrices.
import scipy.linalg

# For creating tensor networks and performing contractions.
import tensornetwork as tn

# Assign an alias to the ``spinbosonchain`` library.
import spinbosonchain as sbc

# For performing SVD truncation sweeps, shifting orthogonal centers,
# and single-node SVD.
import spinbosonchain._svd

# For performing QR sweeps, and single-node QR factorizations.
import spinbosonchain._qr

# For finding the dominant eigenvectors of specially structured tensor networks.
import spinbosonchain._arnoldi

############################
## Authorship information ##
############################

__author__     = "D-Wave Systems Inc."
__copyright__  = "Copyright 2021"
__credits__    = ["Matthew Fitzpatrick"]
__maintainer__ = "D-Wave Systems Inc."
__email__      = "support@dwavesys.com"
__status__     = "Development"

##################################
## Define classes and functions ##
##################################

def apply_mpo_to_mps_and_compress(mpo_nodes,
                                  mps_nodes,
                                  compress_params,
                                  is_infinite):
    kwargs = {"mpo_nodes": mpo_nodes,
              "mps_nodes": mps_nodes,
              "compress_params": compress_params}
    if is_infinite:
        # The Schmidt spectrum for bonds between unit cells is calculated.
        schmidt_spectra = \
            apply_infinite_mpo_to_infinite_mps_and_compress(**kwargs)
    else:
        # The Schmidt spectra is calculated for all bonds if no variational
        # compression is to be performed, otherwise schmidt_spectra is set to
        # `None`.
        schmidt_spectra = \
            apply_finite_mpo_to_finite_mps_and_compress(**kwargs)

    return schmidt_spectra

def apply_infinite_mpo_to_infinite_mps_and_compress(mpo_nodes,
                                                    mps_nodes,
                                                    compress_params):
    apply_directly_mpo_to_mps(mpo_nodes, mps_nodes)

    # See comments in function spinbosonchain._svd.Lambda_Theta_form for a
    # description of the Lambda_Theta object.
    Lambda_Theta = sbc._svd.Lambda_Theta_form(mps_nodes)

    canonicalize_and_compress_infinite_mps(Lambda_Theta, compress_params)

    Lambda = Lambda_Theta[0]
    Theta_nodes = Lambda_Theta[1]

    L = len(mps_nodes)

    nodes_to_contract = [Lambda, Theta_nodes[0]]
    network_struct = [(-1, 1), (1, -2, -3)]
    mps_nodes[0] = tn.ncon(nodes_to_contract, network_struct)

    for r in range(1, L):
        mps_nodes[r] = Theta_nodes[r]

    schmidt_spectra = [Lambda]

    return schmidt_spectra

def apply_finite_mpo_to_finite_mps_and_compress(mpo_nodes,
                                                mps_nodes,
                                                compress_params):
    # [1]: Annals of Physics 326 (2011) 96-192.
    # [2]: New J. Phys. 12, 055026 (2010).

    # Put MPO in left-canonical form. MPS is assumed to be already in said form.
    kwargs = {"nodes": mpo_nodes, "normalize": False}
    sbc._qr.left_to_right_sweep(**kwargs)

    # See docs of class spinbosonchain.compress.Params for descriptions of the
    # various compression parameters used below.
    if compress_params.max_num_var_sweeps > 0:
        norm_of_mps_to_compress = \
            norm_of_mpo_mps_network_wo_compression(mpo_nodes, mps_nodes)
        initial_mps_nodes = copy.deepcopy(mps_nodes)
    else:
        initial_mps_nodes = None

    if compress_params.method == "direct":
        # Here we essentially follow the discussion in the paragraph containing
        # Eq. (16) of [2], and the subsequent paragraph, with a few changes:
        # we first perform a sweep from right to left without compression,
        # followed by a return sweep from left to right. Since there is no
        # compression in first sweep, we can use QR instead of SVD.
        apply_directly_mpo_to_mps(mpo_nodes, mps_nodes)
        kwargs["nodes"] = mps_nodes
        sbc._qr.right_to_left_sweep(**kwargs)
    else:
        # For a discussion on the zip-up method, see the paragraph above that
        # containing Eq. (17) of [2], and read through until the end of that
        # section. Note that mps_nodes is updated in place.
        zip_up(mpo_nodes, mps_nodes, compress_params)

    # Perform a return SVD truncation sweep. If the SVD truncation sweep is
    # followed by variational compression, then L-nodes are calculated and
    # cached while performing the SVD truncation sweep for efficiency. See
    # Sec. 4.5.2 of [1] for a discussion on the L-nodes and variational
    # compression. Note that mps_nodes is updated in place.
    kwargs = {"mpo_nodes": mpo_nodes,
              "mps_nodes": mps_nodes,
              "initial_mps_nodes": initial_mps_nodes,
              "compress_params": compress_params}
    L_cache, schmidt_spectra = return_svd_sweep(**kwargs)

    if compress_params.max_num_var_sweeps > 0:
        # If variational compression is performed, then the Schmidt spectra is
        # not calculated.
        schmidt_spectra = None
        kwargs["norm_of_mps_to_compress"] = norm_of_mps_to_compress
        kwargs["L_cache"] = L_cache
        variational_compression(**kwargs)

    return schmidt_spectra

def norm_of_mpo_mps_network_wo_compression(mpo_nodes, mps_nodes):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # See Sec. 4.5.2 of [1] for some helpful context. In this scenario, the
    # mpo-mps network without compression is denoted by |psi> in Sec. 4.5.2 of
    # [1]. The R-nodes are discussed in Sec. 4.5.2 of [1].
    imax = len(mps_nodes) - 1
    R = trivial_R(mpo_nodes, mps_nodes, num_legs=4)  # Right-most R-node.

    for i in range(imax, -1, -1):
        MWWR = contract_MWWR_network(mps_nodes[i], mpo_nodes[i], R)
        conj_mps_node = tn.conj(mps_nodes[i])
        conj_mps_node[1] ^ MWWR[3]
        conj_mps_node[2] ^ MWWR[4]
        output_edge_order = (MWWR[0], MWWR[1], MWWR[2], conj_mps_node[0])
        R = tn.contract_between(node1=conj_mps_node,
                                node2=MWWR,
                                output_edge_order=output_edge_order)

    nodes_to_contract = [R]
    network_struct = [(1, 2, 2, 1)]
    result = tn.ncon(nodes_to_contract, network_struct)
    result = float(np.sqrt(np.real(result.tensor)))

    return result

def trivial_R(mpo_nodes, mps_nodes, num_legs):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # See Sec. 4.5.2 of [1] for a discussion on the R-nodes. The 'trivial'
    # R-node is essentially a multi-dimensional Kronecker delta.

    w_r = mpo_nodes[-1].shape[-1]  # w_r: right dangling mpo bond dimension.
    chi_r = mps_nodes[-1].shape[-1]  # chi_r: right dangling mps bond dimension.
    if num_legs == 3:
        tensor = np.zeros([chi_r, w_r, chi_r*w_r], dtype=np.complex128)
        for m2 in range(chi_r*w_r):
            m1 = 0 if chi_r == 1 else m2
            w = 0 if w_r == 1 else m2
            tensor[m1, w, m2] = 1
        R = tn.Node(tensor)
    else:
        tensor = np.zeros([chi_r, w_r, w_r, chi_r], dtype=np.complex128)
        for m in range(chi_r):
            for w in range(w_r):
                tensor[m, w, w, m] = 1
        R = tn.Node(tensor)

    return R

def trivial_L(mpo_nodes, mps_nodes):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # See Sec. 4.5.2 of [1] for a discussion on the L-nodes. The 'trivial'
    # L-node is essentially a multi-dimensional Kronecker delta.

    if mps_nodes is None:
        L = None
        return L

    w_l = mpo_nodes[0].shape[0]  # w_l: left dangling mpo bond dimension.
    chi_l = mps_nodes[0].shape[0]  # chi_l: left dangling mps bond dimension.
    tensor = np.zeros([chi_l, w_l, chi_l*w_l], dtype=np.complex128)
    for m2 in range(chi_l*w_l):
        m1 = 0 if chi_l == 1 else m2
        w = 0 if w_l == 1 else m2
        tensor[m1, w, m2] = 1
    L = tn.Node(tensor)

    return L

def contract_MR_network(M, R, num_R_legs=3):
    # See comments in function contract_MWWR_network for context.

    nodes_to_contract = [M, R]
    if num_R_legs == 3:
        network_struct = [(-1, -2, 1), (1, -3, -4)]
    else:
        network_struct = [(-1, -2, 1), (1, -3, -4, -5)]
    MR = tn.ncon(nodes_to_contract, network_struct)

    return MR

def contract_MWR_network(M, W, R, num_R_legs=3):
    # See comments in function contract_MWWR_network for context.

    MR = contract_MR_network(M, R, num_R_legs)

    W[2] ^ MR[1]
    W[3] ^ MR[2]
    if num_R_legs == 3:
        output_edge_order = (MR[0], W[0], W[1], MR[3])

    else:
        output_edge_order = (MR[0], W[0], W[1], MR[3], MR[4])
    MWR = tn.contract_between(node1=W,
                              node2=MR,
                              output_edge_order=output_edge_order)

    return MWR

def contract_MWWR_network(M, W, R):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # M is the MPS node at the current site of interest. W is the MPO node being
    # applied to M. R is the R-node just to the right of the current site of
    # interest. See Sec. 4.5.2 of [1] for a discussion on the R-nodes.

    MWR = contract_MWR_network(M, W, R, num_R_legs=4)
    W = tn.conj(W)

    W[2] ^ MWR[2]
    W[3] ^ MWR[3]
    output_edge_order = (MWR[0], MWR[1], W[0], W[1], MWR[4])
    MWWR = tn.contract_between(node1=W,
                               node2=MWR,
                               output_edge_order=output_edge_order)

    return MWWR

def contract_LM_network(L, M):
    # See comments in function contract_LMW_network for context.

    nodes_to_contract = [L, M]
    network_struct = [(1, -3, -4), (1, -2, -1)]
    LM = tn.ncon(nodes_to_contract, network_struct)

    return LM

def contract_LMW_network(L, M, W):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # M is the MPS node at the current site of interest to which a MPO is to be
    # applied at some point. L is the L-node just to the left of the current
    # site of interest. See Sec. 4.5.2 of [1] for a discussion on the L-nodes.

    LM = contract_LM_network(L, M)

    LM[1] ^ W[2]
    LM[2] ^ W[0]
    output_edge_order = (LM[0], W[3], W[1], LM[3])
    LMW = tn.contract_between(node1=LM,
                              node2=W,
                              output_edge_order=output_edge_order)

    return LMW

def zip_up(mpo_nodes, mps_nodes, compress_params):
    # [2]: New J. Phys. 12, 055026 (2010).

    # For a discussion on the zip-up method, see the paragraph above that
    # containing Eq. (17) of [2], and read through until the end of that
    # section.

    imax = len(mps_nodes) - 1
    U = None
    S = None

    # See second last paragraph of Sec. 2.1.3 of [2] for the reasoning behind
    # the code block below.
    compress_params = copy.deepcopy(compress_params)
    if compress_params.max_num_singular_values is not None:
        compress_params.max_num_singular_values *= 2
    if compress_params.max_trunc_err is not None:
        compress_params.max_trunc_err /= 10
    if compress_params.svd_rel_tol is not None:
        compress_params.svd_rel_tol /= 10

    for i in range(imax, -1, -1):
        node_idx = i
        U, S = update_mps_node_in_zip_up(node_idx,
                                         mpo_nodes,
                                         mps_nodes,
                                         U,
                                         S,
                                         compress_params)

    return None

def update_mps_node_in_zip_up(node_idx,
                              mpo_nodes,
                              mps_nodes,
                              U,
                              S,
                              compress_params):
    # [2]: New J. Phys. 12, 055026 (2010).

    # For a discussion on the zip-up method, see the paragraph above that
    # containing Eq. (17) of [2], and read through until the end of that
    # section.

    i = node_idx
    imax = len(mps_nodes) - 1
    mpo_node = mpo_nodes[i]
    mps_node = mps_nodes[i]

    if i == imax:
        nodes_to_contract = (mpo_node, mps_node)
        network_struct = [(-1, -3, 1, -4), (-2, 1, -5)]
        temp_node_2 = tn.ncon(nodes_to_contract, network_struct)
        tn.flatten_edges([temp_node_2[3], temp_node_2[4]])
    elif 0 <= i < imax:
        nodes_to_contract = (mps_node, U, S)
        network_struct = [(-1, -2, 2), (-3, 2, 1), (1, -4)]
        temp_node_1 = tn.ncon(nodes_to_contract, network_struct)
        temp_node_1[1] ^ mpo_node[2]
        temp_node_1[2] ^ mpo_node[3]
        output_edge_order = \
            (mpo_node[0], temp_node_1[0], mpo_node[1], temp_node_1[3])
        temp_node_2 = tn.contract_between(node1=temp_node_1,
                                          node2=mpo_node,
                                          output_edge_order=output_edge_order)

    if 0 < i <= imax:
        left_edges = (temp_node_2[0], temp_node_2[1])
        right_edges = (temp_node_2[2], temp_node_2[3])
        U, S, V_dagger = sbc._svd.split_node_full(temp_node_2,
                                                  left_edges,
                                                  right_edges,
                                                  compress_params)
        mps_nodes[i] = V_dagger
    elif i == 0:
        tn.flatten_edges([temp_node_2[0], temp_node_2[1]])
        temp_node_2.reorder_edges([temp_node_2[2],
                                   temp_node_2[0],
                                   temp_node_2[1]])
        mps_nodes[i] = temp_node_2

    return U, S

def return_svd_sweep(mpo_nodes, mps_nodes, initial_mps_nodes, compress_params):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # Perform a return SVD truncation sweep. If the SVD truncation sweep is
    # followed by variational compression, then L-nodes are calculated and
    # cached while performing the SVD truncation sweep for efficiency. See
    # Sec. 4.5.2 of [1] for a discussion on the L-nodes and variational
    # compression. Readers should note that the mpo-mps network without
    # compression is denoted by |psi> in Sec. 4.5.2 and should keep this in
    # mind when reading about the L-nodes.

    # Note further that initial_mps_nodes represents the MPS before the MPO
    # application, whereas at the end of this function, mps_nodes represents
    # the MPS obtained after applying the MPO.

    imin = 0
    imax = len(mps_nodes) - 2
    L = trivial_L(mpo_nodes, initial_mps_nodes)
    L_cache = [L]
    schmidt_spectra = []

    # Normalize MPS now if no variational compression is to be performed.
    # Otherwise, normalization is performed after variational compression.
    normalize_schmidt_spectra = (compress_params.max_num_var_sweeps == 0)

    for i in range(imin, imax+1):
        kwargs = {"nodes": mps_nodes,
                  "current_orthogonal_center_idx": i,
                  "compress_params": compress_params,
                  "normalize_schmidt_spectra": normalize_schmidt_spectra}
        U, S, V_dagger = sbc._svd.shift_orthogonal_center_to_the_right(**kwargs)
        schmidt_spectra.insert(0, S)

        if compress_params.max_num_var_sweeps > 0:
            M = initial_mps_nodes[i]
            W = mpo_nodes[i]
            LMW = contract_LMW_network(L, M, W)
            conj_mps_node = tn.conj(mps_nodes[i])
            LMW[2] ^ conj_mps_node[1]
            LMW[3] ^ conj_mps_node[0]
            output_edge_order = (LMW[0], LMW[1], conj_mps_node[2])
            L = tn.contract_between(node1=LMW,
                                    node2=conj_mps_node,
                                    output_edge_order=output_edge_order)
            L_cache.append(L)

    if compress_params.max_num_var_sweeps == 0:
        # Normalize MPS such that its 'norm' equal unity.
        mps_nodes[-1] /= tn.norm(mps_nodes[-1])

    return L_cache, schmidt_spectra

def variational_compression(mpo_nodes,
                            mps_nodes,
                            initial_mps_nodes,
                            norm_of_mps_to_compress,
                            L_cache,
                            compress_params):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # This function essentially implements the variational compression procedure
    # described in Sec. 4.5.2 of [1], see therein for discussion on the L- and
    # R-nodes. Readers should note that the mpo-mps network without compression
    # is denoted by |psi> in Sec. 4.5.2 and should keep this in mind when
    # reading about the L- and R-nodes.

    # initial_mps_nodes represents the MPS before the MPO application, whereas
    # mps_nodes represents the MPS obtained after applying the MPO.

    imax = len(mps_nodes) - 1
    R_cache = [trivial_R(mpo_nodes, initial_mps_nodes, num_legs=3)]

    # A full sweep goes right to left, then left to right.
    sweep_count = 0
    while sweep_count < compress_params.max_num_var_sweeps:
        kwargs = {"mpo_nodes": mpo_nodes,
                  "mps_nodes": mps_nodes,
                  "initial_mps_nodes": initial_mps_nodes,
                  "norm_of_mps_to_compress": norm_of_mps_to_compress,
                  "L_cache": L_cache,
                  "R_cache": R_cache,
                  "compress_params": compress_params}

        if variational_compression_has_converged(**kwargs):
            break

        kwargs = {"mpo_nodes": mpo_nodes,
                  "mps_nodes": mps_nodes,
                  "initial_mps_nodes": initial_mps_nodes,
                  "L_cache": L_cache,
                  "R_cache": R_cache}

        for i in range(imax, 0, -1):
            update_node_and_shift_left_in_variational_compression(**kwargs)
        for i in range(imax):
            update_node_and_shift_right_in_variational_compression(**kwargs)

        sweep_count += 1

    # Normalize MPS such that its 'norm' equal unity.
    mps_nodes[-1] /= tn.norm(mps_nodes[-1])

    return None

def variational_compression_has_converged(mpo_nodes,
                                          mps_nodes,
                                          initial_mps_nodes,
                                          norm_of_mps_to_compress,
                                          L_cache,
                                          R_cache,
                                          compress_params):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # See Sec. 4.5.2 of [1] for a discussion on the variational compression
    # procedure implemented in this module. Readers should note that the
    # mpo-mps network without compression is denoted by |psi> in Sec. 4.5.2.

    # This function checks whether the variational compression procedure has
    # achieved convergence. This is done by evaluating the left hand side of
    # Eq. (150) of [1].

    # 'compressed' and 'uncompressed' MPS's are refering to |psi'> and |psi>
    # in Sec. 4.5.2 [see above note on |psi>]. Hence, in both cases these are
    # MPS's obtained after applying the MPO.

    # initial_mps_nodes represents the MPS before the MPO application, whereas
    # mps_nodes represents the MPS obtained after applying the MPO.

    overlap = overlap_btwn_compressed_and_uncompressed_mps(L_cache,
                                                           initial_mps_nodes,
                                                           mpo_nodes,
                                                           mps_nodes,
                                                           R_cache)

    nodes_to_contract = (mps_nodes[-1], tn.conj(mps_nodes[-1]))
    network_struct = [(1, 2, 3), (1, 2, 3)]
    norm_sq_of_compressed_mps = tn.ncon(nodes_to_contract, network_struct)
    norm_of_compressed_mps = \
        float(np.sqrt(np.real(norm_sq_of_compressed_mps.tensor)))

    var_rel_err_sq = (np.abs(norm_of_mps_to_compress * norm_of_mps_to_compress
                             - 2 * np.real(overlap)
                             + norm_of_compressed_mps * norm_of_compressed_mps)
                      / norm_of_mps_to_compress)
    var_rel_err = np.sqrt(var_rel_err_sq)

    if var_rel_err < compress_params.var_rel_tol:
        result = True
    else:
        result = False

    return result

def overlap_btwn_compressed_and_uncompressed_mps(L_cache,
                                                 initial_mps_nodes,
                                                 mpo_nodes,
                                                 mps_nodes,
                                                 R_cache):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # See Sec. 4.5.2 of [1] for a discussion on the variational compression
    # procedure implemented in this module. Readers should note that the
    # mpo-mps network without compression is denoted by |psi> in Sec. 4.5.2.

    # 'compressed' and 'uncompressed' MPS's are refering to |psi'> and |psi>
    # in Sec. 4.5.2 [see above note on |psi>]. Hence, in both cases these are
    # MPS's obtained after applying the MPO.

    # initial_mps_nodes represents the MPS before the MPO application, whereas
    # mps_nodes represents the MPS obtained after applying the MPO.

    L = L_cache[-1]
    M = initial_mps_nodes[-1]
    W = mpo_nodes[-1]
    R = R_cache[0]

    MWR = contract_MWR_network(M, W, R)

    L[0] ^ MWR[0]
    L[1] ^ MWR[1]
    output_edge_order = (L[2], MWR[2], MWR[3])
    LMWR = tn.contract_between(node1=L,
                               node2=MWR,
                               output_edge_order=output_edge_order)

    conj_mps_node = tn.conj(mps_nodes[-1])

    for edge_idx in range(3):
        conj_mps_node[edge_idx] ^ LMWR[edge_idx]

    overlap = tn.contract_between(node1=conj_mps_node, node2=LMWR)
    overlap = complex(overlap.tensor)

    return overlap

def update_node_and_shift_left_in_variational_compression(mpo_nodes,
                                                          mps_nodes,
                                                          initial_mps_nodes,
                                                          L_cache,
                                                          R_cache):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # This function essentially implements a single shift to the left in the
    # variational compression procedure described in Sec. 4.5.2 of [1], see
    # therein for discussion on the L- and R-nodes. Readers should note that the
    # mpo-mps network without compression is denoted by |psi> in Sec. 4.5.2 and
    # should keep this in mind when reading about the L- and R-nodes.

    # initial_mps_nodes represents the MPS before the MPO application, whereas
    # mps_nodes represents the MPS obtained after applying the MPO.

    orthogonal_center_idx = len(L_cache) - 1
    L = L_cache[-1]
    M = initial_mps_nodes[orthogonal_center_idx]
    W = mpo_nodes[orthogonal_center_idx]
    R = R_cache[0]

    MWR = contract_MWR_network(M, W, R)

    L[0] ^ MWR[0]
    L[1] ^ MWR[1]
    output_edge_order = (L[2], MWR[2], MWR[3])
    LMWR = tn.contract_between(node1=L,
                               node2=MWR,
                               output_edge_order=output_edge_order)

    L_cache.pop()

    QR = tn.conj(LMWR)
    left_node, right_node = sbc._qr.split_node(node=QR,
                                               left_edges=(QR[1], QR[2]),
                                               right_edges=(QR[0],))
    new_edge_order = [left_node[2], left_node[0], left_node[1]]
    B = tn.conj(left_node.reorder_edges(new_edge_order))
    mps_nodes[orthogonal_center_idx] = B
    conj_B = tn.conj(B)

    conj_B[1] ^ MWR[2]
    conj_B[2] ^ MWR[3]
    output_edge_order = (MWR[0], MWR[1], conj_B[0])
    next_R = tn.contract_between(node1=conj_B,
                                 node2=MWR,
                                 output_edge_order=output_edge_order)
    R_cache.insert(0, next_R)

    return None

def update_node_and_shift_right_in_variational_compression(mpo_nodes,
                                                           mps_nodes,
                                                           initial_mps_nodes,
                                                           L_cache,
                                                           R_cache):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # This function essentially implements a single shift to the right in the
    # variational compression procedure described in Sec. 4.5.2 of [1], see
    # therein for discussion on the L- and R-nodes. Readers should note that the
    # mpo-mps network without compression is denoted by |psi> in Sec. 4.5.2 and
    # should keep this in mind when reading about the L- and R-nodes.

    # initial_mps_nodes represents the MPS before the MPO application, whereas
    # mps_nodes represents the MPS obtained after applying the MPO.

    orthogonal_center_idx = len(L_cache) - 1
    L = L_cache[-1]
    M = initial_mps_nodes[orthogonal_center_idx]
    W = mpo_nodes[orthogonal_center_idx]
    R = R_cache[0]

    LMW = contract_LMW_network(L, M, W)

    LMW[0] ^ R[0]
    LMW[1] ^ R[1]
    output_edge_order = (LMW[3], LMW[2], R[2])
    LMWR = tn.contract_between(node1=LMW,
                               node2=R,
                               output_edge_order=output_edge_order)
    R_cache.pop(0)

    QR = LMWR
    left_node, right_node = sbc._qr.split_node(node=QR,
                                               left_edges=(QR[0], QR[1]),
                                               right_edges=(QR[2],))
    A = left_node
    mps_nodes[orthogonal_center_idx] = A
    conj_A = tn.conj(A)

    if orthogonal_center_idx == len(mps_nodes) - 2:
        nodes_to_contract = [right_node, mps_nodes[-1]]
        network_struct = [(-1, 1), (1, -2, -3)]
        mps_nodes[-1] = tn.ncon(nodes_to_contract, network_struct)

    LMW[2] ^ conj_A[1]
    LMW[3] ^ conj_A[0]
    output_edge_order = (LMW[0], LMW[1], conj_A[2])
    next_L = tn.contract_between(node1=LMW,
                                 node2=conj_A,
                                 output_edge_order=output_edge_order)
    L_cache.append(next_L)

    return None

def apply_directly_mpo_to_mps(mpo_nodes, mps_nodes):
    for idx, (mpo_node, mps_node) in enumerate(zip(mpo_nodes, mps_nodes)):
        mps_nodes[idx] = apply_directly_mpo_node_to_mps_node(mpo_node, mps_node)

    return None

def apply_directly_mpo_node_to_mps_node(mpo_node, mps_node):
    nodes_to_contract = (mpo_node, mps_node)
    network_struct = [(-1, -3, 1, -4), (-2, 1, -5)]
    new_mps_node = tn.ncon(nodes_to_contract, network_struct)

    # A single call to tn.flatten_edges will take the flatten edge set and move
    # it to the right end of the list of edges, hence the necessary call to
    # reorder_edges.
    tn.flatten_edges([new_mps_node[0], new_mps_node[1]])
    tn.flatten_edges([new_mps_node[1], new_mps_node[2]])
    new_mps_node.reorder_edges([new_mps_node[1],
                                new_mps_node[0],
                                new_mps_node[2]])

    return new_mps_node

def canonicalize_and_compress_infinite_mps(Lambda_Theta, compress_params):
    # [1]: Annals of Physics 326 (2011) 96-192.
    # [3]: Phys. Rev. B 78, 155117 (2008)
    # [4]: arXiv.0804.2509 (2008)

    # This function compresses a given infinite MPS. The procedure implemented
    # here draws from [1], [3], and [4]. Notation-wise we follow mostly [1].
    # In [1], the relevant content is found in Sec. 10.5.

    # See comments in function spinbosonchain._svd.Lambda_Theta_form for a
    # description of the Lambda_Theta object.

    # V_L, V_R, X, X_inv, Y, and Y_inv are the same as those defined in
    # Sec. 10.5 of [1] except that we have generalized to a L-site unit cell
    # where L>0.
    V_L = calc_V_L(Lambda_Theta)
    V_R = calc_V_R(Lambda_Theta)
    X, X_inv = calc_X_and_X_inv(V_L)
    Y, Y_inv = calc_Y_and_Y_inv(V_R)

    old_Lambda = Lambda_Theta[0]
    old_Theta_nodes = Lambda_Theta[1]

    # Here we obtain the new singular value matrix for the bond between unit
    # cells. Note that the singular value spectrum is truncated.
    X_old_Lambda_Y = tn.ncon([X, old_Lambda, Y], [(-1, 1), (1, 2), (2, -2)])
    U, S, V_dagger = sbc._svd.split_node_full(node=X_old_Lambda_Y,
                                              left_edges=(X_old_Lambda_Y[0],),
                                              right_edges=(X_old_Lambda_Y[1],),
                                              compress_params=compress_params)

    new_Lambda = S / tn.norm(S)

    L = len(old_Theta_nodes)

    # Here we construct the Z_L and Z_R that are used to transform the Theta
    # node, which is a part of the canonicalization procedure.
    if L == 1:
        # Here we are essentially following Fig. 2.ii of [3]. Note that our
        # V_dagger is their V; our Y_inv is their X_inv; and our X_inv is their
        # Y_T_inv.
        Z_L = tn.ncon([V_dagger, Y_inv], [(-1, 1), (1, -2)])
        Z_R = tn.ncon([X_inv, U], [(-1, 1), (1, -2)])
    else:
        # Here we are essentially following the paragraph containing Eq. (355)
        # of [1] with a few minor changes.
        new_Lambda_inv_tensor = \
            np.diag(1 / np.diag(np.array(new_Lambda.tensor)))
        new_Lambda_inv = tn.Node(new_Lambda_inv_tensor)

        conj_U = tn.conj(U)
        V_T = tn.conj(V_dagger)

        nodes_to_contract = [new_Lambda_inv, conj_U, X, old_Lambda]
        network_struct = [(-1, 3), (1, 3), (1, 2), (2, -2)]
        Z_L = tn.ncon(nodes_to_contract, network_struct)

        nodes_to_contract = [old_Lambda, Y, V_T, new_Lambda_inv]
        network_struct = [(-1, 1), (1, 2), (3, 2), (3, -2)]
        Z_R = tn.ncon(nodes_to_contract, network_struct)

    new_Theta_nodes = old_Theta_nodes

    nodes_to_contract = [Z_L, new_Theta_nodes[0]]
    network_struct = [(-1, 1), (1, -2, -3)]
    new_Theta_nodes[0] = tn.ncon(nodes_to_contract, network_struct)

    nodes_to_contract = [new_Theta_nodes[-1], Z_R]
    network_struct = [(-1, -2, 1), (1, -3)]
    new_Theta_nodes[-1] = tn.ncon(nodes_to_contract, network_struct)

    compress_and_normalize_Theta(new_Lambda, new_Theta_nodes, compress_params)

    Lambda_Theta[0] = new_Lambda
    Lambda_Theta[1] = new_Theta_nodes

    return None

def compress_and_normalize_Theta(Lambda, Theta_nodes, compress_params):
    # [3]: Phys. Rev. B 78, 155117 (2008)

    # See comments in function spinbosonchain._svd.Lambda_Theta_form for a
    # description of the Lambda_Theta object.

    # Essentially following [3], here we compress and normalize the Theta node
    # of the Lambda-Theta pair. We apply a direct SVD compression procedure.

    kwargs = {"nodes": Theta_nodes, "normalize": False}
    sbc._qr.left_to_right_sweep(**kwargs)
    kwargs["compress_params"] = compress_params
    kwargs["normalize"] = True
    sbc._svd.right_to_left_sweep(**kwargs)

    # Normalize the Theta node such that the MPS it represents has a state
    # 'norm' equal to unity.
    M = tn.ncon([Lambda, Theta_nodes[0]], [(-1, 1), (1, -2, -3)])
    Theta_nodes[0] /= tn.norm(M)

    return None

def calc_V_L(Lambda_Theta):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # V_L is the same as that defined in Sec. 10.5 of [1] except that we have
    # generalized to a L-site unit cell where L>0.

    kwargs = {"Lambda_Theta": Lambda_Theta,
              "unit_cell_type": "left",
              "krylov_dim": 10,
              "epsilon_D": 1.0e-14}

    mu_L, V_L = sbc._arnoldi.dominant_eigpair_of_transfer_matrix(**kwargs)

    return V_L

def calc_V_R(Lambda_Theta):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # V_R is the same as that defined in Sec. 10.5 of [1] except that we have
    # generalized to a L-site unit cell where L>0.

    kwargs = {"Lambda_Theta": Lambda_Theta,
              "unit_cell_type": "right",
              "krylov_dim": 10,
              "epsilon_D": 1.0e-14}

    mu_R, V_R = sbc._arnoldi.dominant_eigpair_of_transfer_matrix(**kwargs)

    return V_R

def calc_X_and_X_inv(V_L):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # V_L, X, and X_inv are the same as those defined in Sec. 10.5 of [1] except
    # that we have generalized to a L-site unit cell where L>0.

    if V_L.backend.name != "numpy":
        sbc._backend.tf_to_np(V_L)

    try:
        D, W = scipy.linalg.eigh(V_L.tensor, driver='ev')
    except TypeError as err:  # Occurs with some versions scipy.
        D, W = scipy.linalg.eigh(V_L.tensor)
    D = np.abs(D)
    tol = 1.0e-14
    D[D<tol] = tol

    X_dagger = W @ np.diag(np.sqrt(D))
    X = np.conj(np.transpose(X_dagger))
    X = tn.Node(X)

    # Note that X_inv is really a pseudo-inverse of X.
    X_inv = W @ np.diag(1/np.sqrt(D))
    X_inv = tn.Node(X_inv)

    return X, X_inv

def calc_Y_and_Y_inv(V_R):
    # [1]: Annals of Physics 326 (2011) 96-192.

    # V_R, Y, and Y_inv are the same as those defined in Sec. 10.5 of [1] except
    # that we have generalized to a L-site unit cell where L>0.

    if V_R.backend.name != "numpy":
        sbc._backend.tf_to_np(V_R)

    try:
        D, W = scipy.linalg.eigh(V_R.tensor, driver='ev')
    except TypeError as err:  # Occurs with some versions scipy.
        D, W = scipy.linalg.eigh(V_R.tensor)
    D = np.abs(D)
    tol = 1.0e-14
    D[D<tol] = tol

    Y = W @ np.diag(np.sqrt(D))
    Y = tn.Node(Y)  # Convert from numpy array to tensornetwork node.

    # Note that Y_inv is really a pseudo-inverse of Y.
    W_dagger = np.conj(np.transpose(W))
    Y_inv = np.diag(1/np.sqrt(D)) @ W_dagger
    Y_inv = tn.Node(Y_inv)

    return Y, Y_inv