import contextlib
import os
import pathlib
import tempfile
from dataclasses import dataclass
from typing import Optional, Union

import isce3
import numpy as np
from isce3.ext.isce3.unwrap import _snaphu_unwrap

@dataclass(frozen=True)
class TopoCostParams:
    r"""Configuration parameters for SNAPHU "topo" cost mode

    Unwrapped phase is modelled as a topographic SAR interferometry signal
    in additive phase noise.

    Backscatter brightness is modelled according to the scattering model:

    .. math::

        \sigma^0 = C * \left( kds * \cos \theta_{inc} +
                \cos ^n (2 \theta_{inc}) \right) * \cos \theta_{inc}

    where :math:`C` is a scaling factor, :math:`kds` is the ratio of diffuse to
    specular scattering, :math:`\theta_{inc}` is the local incidence angle, and
    :math:`n` (input parameter `specular_exp`) is the power to which the
    specular cosine term is raised.

    Attributes
    ----------
    bperp : float
        Perpendicular baseline length, in meters. If the value is negative,
        increasing phase implies increasing topographic height.
    near_range : float
        Slant range from platform to the first range bin, in meters.
    dr, da : float
        Range & azimuth bin spacing after any multi-looking, in meters.
    range_res, az_res : float
        Single-look range & azimuth resolution, in meters.
    wavelength : float
        Wavelength, in meters.
    transmit_mode : {"pingpong", "repeat_pass", "single_antenna_transmit"}
        Radar transmit mode. 'pingpong' and 'repeat_pass' modes indicate that
        both antennas both transmitted and received. Both modes have the same
        effect in the algorithm. 'single_antenna_transmit' indicates that a
        single antenna was used to transmit while both antennas received. In
        this mode, the baseline is effectively halved.
    altitude : float
        Platform altitude relative to the Earth's surface, in meters.
    earth_radius : float, optional
        Local Earth radius, in meters. A spherical-Earth model is used.
        (default: 6378000.0)
    kds : float, optional
        Ratio of diffuse to specular scattering. (default: 0.02)
    specular_exp : float, optional
        Power specular scattering component. Larger values imply a sharper peak
        for specular scatter. (default: 8.0)
    dzr_crit_factor : float, optional
        Multiplicative factor applied to diffuse scatter term in evaluating
        crossover point between diffuse and specular scatter in terms of range
        slope. (default: 2.0)
    shadow : bool, optional
        Allow discontinuities from shadowing? If this is disabled, the minimum
        topographic slope estimated from mean backscatter intensity is clipped
        to the value of `dz_ei_min`. (default: False)
    dz_ei_min : float, optional
        Minimum slope expected in the absence of layover, in meters per
        slant-range pixel. (default: -4.0)
    lay_width : int, optional
        Width of window (number of pixels) for summing layover brightness.
        (default: 16)
    lay_min_ei : float, optional
        Threshold brightness (normalized) for assuming layover. (default: 1.25)
    slope_ratio_factor : float, optional
        Multiplicative factor applied to kds in order to get ratio of slopes for
        linearized scattering model. The term improves agreement of the
        piecewise-linear model with the cosine model near the transition point
        (dzrcrit) at the expense of poorer agreement at very large slopes.
        (default: 1.18)
    sigsq_ei : float, optional
        Variance of range slopes due to uncertainties in slope estimation from
        brightness, in (meters/pixel)^2. (default: 100.0)
    drho : float, optional
        Step size for calculating lookup table of maximum layover slope based on
        measured correlation. (default: 0.005)
    dz_lay_peak : float, optional
        Layover peak location, in meters/pixel. (default: -2.0)
    azdz_factor : float, optional
        Factor applied to range layover probability density to get azimuth
        layover probability density. (default: 0.99)
    dz_ei_factor : float, optional
        Factor applied to slope expected from brightness without layover. Can
        account for underestimation of brightness from averaging with
        neighboring dark pixels when despeckling. (default: 4.0)
    dz_ei_weight : float, optional
        Weight applied to slope expected from brightness without layover. Must
        be between zero and one. Can reduce influence of intensity on
        non-layover slope. This is useful if there are lots of non-topographic
        variations in brightness (i.e. changes in surface reflectivity).
        (default: 0.5)
    dz_lay_factor : float, optional
        Factor applied to slope expected from brightness with layover. Can
        account for underestimation of brightness from averaging with
        neighboring dark pixels when despeckling. (default: 1.0)
    lay_const : float, optional
        Ratio of layover probability density to peak probability density for
        non-layover slopes expected. (default: 0.9)
    lay_falloff_const : float, optional
        Factor applied to slope variance for non-layover to get falloff of
        probability density after the  upper layover slope limit has been
        exceeded. (default: 2.0)
    sigsq_lay_factor : float, optional
        Fraction of (ambiguity height)^2 to use for slope variance in the
        presence of layover. (default: 0.1)
    krow_ei, kcol_ei : int, optional
        Number of rows & columns to use in sliding average window used for
        normalizing intensity values. (default: 65, 257)
    init_dzr : float, optional
        Initial value of range slope for dzrcrit numerical solution, in
        meters/pixel. (default: 2048.0)
    init_dz_step : float, optional
        Initial range slope step size in dzrhomax numerical solution, in
        meters/pixel. (default: 100.0)
    cost_scale_ambig_ht : float, optional
        Ambiguity height for auto-scaling the `SolverParams.cost_scale`
        parameter to equal 100. This is the amount of height change, in meters,
        that results in a :math:`2 \pi` change in the interferometric phase. The
        cost scale is automatically adjusted to be inversely proportional to the
        midswath ambiguity height. (default: 80.0)
    dnom_inc_angle : float, optional
        Step size, in radians, for dzrhomax lookup table. The index is on the
        flat-earth incidence angle; this is the sample spacing in the table.
        (default: 0.01)
    kpar_dpsi, kperp_dpsi : int, optional
        Number of pixels in sliding window used for averaging wrapped phase
        gradients to get mean non-layover slope, in directions parallel and
        perpendicular to the examined phase difference. (default: 7, 7)
    """

    bperp: float
    near_range: float
    dr: float
    da: float
    range_res: float
    az_res: float
    wavelength: float
    transmit_mode: str
    altitude: float
    earth_radius: float = 6_378_000.0
    kds: float = 0.02
    specular_exp: float = 8.0
    dzr_crit_factor: float = 2.0
    shadow: bool = False
    dz_ei_min: float = -4.0
    lay_width: int = 16
    lay_min_ei: float = 1.25
    slope_ratio_factor: float = 1.18
    sigsq_ei: float = 100.0
    drho: float = 0.005
    dz_lay_peak: float = -2.0
    azdz_factor: float = 0.99
    dz_ei_factor: float = 4.0
    dz_ei_weight: float = 0.5
    dz_lay_factor: float = 1.0
    lay_const: float = 0.9
    lay_falloff_const: float = 2.0
    sigsq_lay_factor: float = 0.1
    krow_ei: int = 65
    kcol_ei: int = 257
    init_dzr: float = 2048.0
    init_dz_step: float = 100.0
    cost_scale_ambig_ht: float = 80.0
    dnom_inc_angle: float = 0.01
    kpar_dpsi: int = 7
    kperp_dpsi: int = 7

    def tostring(self):
        """Convert to string in SNAPHU config file format."""

        def parse_transmit_mode(s):
            if s == "pingpong":
                return "PINGPONG"
            if s == "repeat_pass":
                return "REPEATPASS"
            if s == "single_antenna_transmit":
                return "SINGLEANTENNATRANSMIT"
            raise ValueError(f"invalid transmit mode '{s}'")

        s = ""
        s += f"BPERP {self.bperp}\n"
        s += f"NEARRANGE {self.near_range}\n"
        s += f"DR {self.dr}\n"
        s += f"DA {self.da}\n"
        s += f"RANGERES {self.range_res}\n"
        s += f"AZRES {self.az_res}\n"
        s += f"LAMBDA {self.wavelength}\n"
        s += f"TRANSMITMODE {parse_transmit_mode(self.transmit_mode)}\n"
        s += f"ALTITUDE {self.altitude}\n"
        s += f"EARTHRADIUS {self.earth_radius}\n"
        s += f"KDS {self.kds}\n"
        s += f"SPECULAREXP {self.specular_exp}\n"
        s += f"DZRCRITFACTOR {self.dzr_crit_factor}\n"
        s += f"SHADOW {self.shadow}\n"
        s += f"DZEIMIN {self.dz_ei_min}\n"
        s += f"LAYWIDTH {self.lay_width}\n"
        s += f"LAYMINEI {self.lay_min_ei}\n"
        s += f"SLOPERATIOFACTOR {self.slope_ratio_factor}\n"
        s += f"SIGSQEI {self.sigsq_ei}\n"
        s += f"DRHO {self.drho}\n"
        s += f"DZLAYPEAK {self.dz_lay_peak}\n"
        s += f"AZDZFACTOR {self.azdz_factor}\n"
        s += f"DZEIFACTOR {self.dz_ei_factor}\n"
        s += f"DZEIWEIGHT {self.dz_ei_weight}\n"
        s += f"DZLAYFACTOR {self.dz_lay_factor}\n"
        s += f"LAYCONST {self.lay_const}\n"
        s += f"LAYFALLOFFCONST {self.lay_falloff_const}\n"
        s += f"SIGSQLAYFACTOR {self.sigsq_lay_factor}\n"
        s += f"KROWEI {self.krow_ei}\n"
        s += f"KCOLEI {self.kcol_ei}\n"
        s += f"INITDZR {self.init_dzr}\n"
        s += f"INITDZSTEP {self.init_dz_step}\n"
        s += f"COSTSCALEAMBIGHT {self.cost_scale_ambig_ht}\n"
        s += f"DNOMINCANGLE {self.dnom_inc_angle}\n"
        s += f"KPARDPSI {self.kpar_dpsi}\n"
        s += f"KPERPDPSI {self.kperp_dpsi}\n"
        return s

@dataclass(frozen=True)
class DefoCostParams:
    """Configuration parameters for SNAPHU "defo" cost mode

    Unwrapped phase is modelled as a surface deformation signature in additive
    phase noise.

    Attributes
    ----------
    azdz_factor : float, optional
        Factor applied to range discontinuity probability density to get
        corresponding value for azimuth. (default: 1.0)
    defo_max : float, optional
        Maximum phase discontinuity, in units of cycles (of 2*pi). If abrupt
        phase discontinuities are not expected, this parameter can be set to
        zero. (default: 1.2)
    sigsq_corr : float, optional
        Phase variance, in cycles^2, reflecting uncertainty in measurement of
        actual statistical correlation. (default: 0.05)
    defo_const : float, optional
        Ratio of phase discontinuity probability density to peak probability
        density expected for discontinuity-possible pixel differences. A value
        of 1 means zero cost for discontinuity, 0 means infinite cost.
        (default: 0.9)
    lay_falloff_const : float, optional
        Factor applied to slope variance for non-layover to get falloff of
        probability density after the  upper layover slope limit has been
        exceeded. (default: 2.0)
    kpar_dpsi, kperp_dpsi : int, optional
        Number of pixels in sliding window used for averaging wrapped phase
        gradients to get mean non-layover slope, in directions parallel and
        perpendicular to the examined phase difference. (default: 7, 7)
    """

    azdz_factor: float = 1.0
    defo_max: float = 1.2
    sigsq_corr: float = 0.05
    defo_const: float = 0.9
    lay_falloff_const: float = 2.0
    kpar_dpsi: int = 7
    kperp_dpsi: int = 7

    def tostring(self):
        """Convert to string in SNAPHU config file format."""
        s = ""
        s += f"DEFOAZDZFACTOR {self.azdz_factor}\n"
        s += f"DEFOMAX_CYCLE {self.defo_max}\n"
        s += f"SIGSQCORR {self.sigsq_corr}\n"
        s += f"DEFOCONST {self.defo_const}\n"
        s += f"LAYFALLOFFCONST {self.lay_falloff_const}\n"
        s += f"KPARDPSI {self.kpar_dpsi}\n"
        s += f"KPERPDPSI {self.kperp_dpsi}\n"
        return s

@dataclass(frozen=True)
class SmoothCostParams:
    """Configuration parameters for SNAPHU "smooth" cost mode

    Attributes
    ----------
    kpar_dpsi, kperp_dpsi : int, optional
        Number of pixels in sliding window used for averaging wrapped phase
        gradients to get mean non-layover slope, in directions parallel and
        perpendicular to the examined phase difference. (default: 7, 7)
    """

    kpar_dpsi: int = 7
    kperp_dpsi: int = 7

    def tostring(self):
        """Convert to string in SNAPHU config file format."""
        s = ""
        s += f"KPARDPSI {self.kpar_dpsi}\n"
        s += f"KPERPDPSI {self.kperp_dpsi}\n"
        return s

@dataclass(frozen=True)
class PNormCostParams:
    r"""Configuration parameters for SNAPHU "p-norm" cost mode

    In this mode, the minimization objective is the :math:`L^p` norm of the
    difference between the unwrapped and wrapped phase gradients.

    .. math:: cost = \sum_i \left| \Delta \phi_i - \Delta \psi_i \right| ^ p

    Attributes
    ----------
    p : float, optional
        Lp norm exponent. Must be nonnegative. (default: 0.0)
    bidir : bool, optional
        If True, bidirectional Lp costs are used. This implies that the scalar
        weight of an Lp arc may be different depending on the direction of net
        flow on the arc. If False, the weight is the same regardless of arc
        direction. (default: True)
    """

    p: float = 0.0
    bidir: bool = True

    def tostring(self):
        """Convert to string in SNAPHU config file format."""
        s = ""
        s += f"PLPN {self.p}\n"
        s += f"BIDIRLPN {self.bidir}\n"
        return s

@dataclass(frozen=True)
class TilingParams:
    """Configuration parameters affecting scene tiling and parallel processing.

    Attributes
    ----------
    nproc : int, optional
        Maximum number of child processes to spawn for parallel tile unwrapping.
        If nproc is less than 1, use all available processors. (default: 1)
    tile_nrows, tile_ncols : int, optional
        Number of tiles along the row/column directions. If `tile_nrows` and
        `tile_ncols` are both 1, the interferogram is unwrapped as a single
        tile. (default: 1, 1)
    row_overlap, col_overlap : int, optional
        Overlap, in number of rows/columns, between neighboring tiles.
        (default: 0)
    tile_cost_thresh : int, optional
        Cost threshold to use for determining boundaries of reliable regions.
        Larger cost threshold implies smaller regions (safer, but more expensive
        computationally). (default: 500)
    min_region_size : int, optional
        Minimum size of a reliable region in tile mode, in pixels. (default: 100)
    tile_edge_weight : float, optional
        Extra weight applied to secondary arcs on tile edges. (default: 2.5)
    secondary_arc_flow_max : int, optional
        Maximum flow magnitude whose cost will be stored in the secondary cost
        lookup table. Secondary costs larger than this will be approximated by a
        quadratic function. (default: 8)
    single_tile_reoptimize : bool, optional
        If True, re-optimize as a single tile after using tile mode for
        initialization. This is equivalent to unwrapping with multiple tiles,
        then using the unwrapped output as the input to a new, single-tile run
        of snaphu to make iterative improvements to the solution. This may
        improve speed compared to a single single-tile run. (default: False)
    """

    nproc: int = 1
    tile_nrows: int = 1
    tile_ncols: int = 1
    row_overlap: int = 0
    col_overlap: int = 0
    tile_cost_thresh: int = 500
    min_region_size: int = 100
    tile_edge_weight: float = 2.5
    secondary_arc_flow_max: int = 8
    single_tile_reoptimize: bool = False

    def tostring(self):
        """Convert to string in SNAPHU config file format."""
        s = ""
        if self.nproc < 1:
            nproc = os.cpu_count() or 1
            s += f"NPROC {nproc}\n"
        else:
            s += f"NPROC {self.nproc}\n"
        s += f"NTILEROW {self.tile_nrows}\n"
        s += f"NTILECOL {self.tile_ncols}\n"
        s += f"ROWOVRLP {self.row_overlap}\n"
        s += f"COLOVRLP {self.col_overlap}\n"
        s += f"TILECOSTTHRESH {self.tile_cost_thresh}\n"
        s += f"MINREGIONSIZE {self.min_region_size}\n"
        s += f"TILEEDGEWEIGHT {self.tile_edge_weight}\n"
        s += f"SCNDRYARCFLOWMAX {self.secondary_arc_flow_max}\n"
        s += f"SINGLETILEREOPTIMIZE {self.single_tile_reoptimize}\n"

        # Don't remove temporary files for each tile since they may be useful
        # for debugging. If the scratch directory is cleaned up, they'll be
        # removed as well.
        s += "RMTMPTILE FALSE\n"

        return s

@dataclass(frozen=True)
class SolverParams:
    """Configuration parameters used by the network initialization and nonlinear
    network flow solver algorithms.

    Attributes
    ----------
    max_flow_inc : int, optional
        Maximum flow increment. (default: 4)
    init_max_flow : int, optional
        Maximum flow to allow in initialization. If this is zero, then the
        maximum is calculated automatically from the statistical cost functions.
        To disable, set it to a large value like 9999, but do not overflow the
        long integer data type. (default: 9999)
    arc_max_flow_const : int, optional
        Constant to add to maximum flow expected from statistical cost functions
        for automatically determining initial maximum flow. (default: 3)
    threshold : float, optional
        Threshold precision for iterative numerical calculations.
        (default: 0.001)
    max_cost : float, optional
        Maximum cost allowed for scalar MST costs and for estimating the number
        of buckets needed for the solver routine. (default: 1000.0)
    cost_scale : float, optional
        Cost scaling factor applied to floating-point costs before quantization
        to integer costs. (default: 100.0)
    n_cycle : int, optional
        Integer spacing that represents one unit of flow (one cycle of phase)
        when storing costs as short integers. (default: 200)
    max_new_node_const : float, optional
        Fraction of total number of nodes to add in each tree expansion phase of
        the solver algorithm. (default: 0.0008)
    max_n_flow_cycles : float or None, optional
        Number of cycles to allow for a call to the solver with a specific flow
        increment delta and still consider that increment done. Ideally it would
        be zero, but scaling for different deltas may leave some negative cycles
        that won't affect the solution much. If None, this is automatically
        determined based on the size of the interferogram. (default: None)
    max_cycle_frac : float, optional
        Fraction of the number of pixels to use as the maximum number of cycles
        allowed for a specific flow increment if max_n_flow_cycles was None.
        (default: 0.00001)
    n_conn_node_min : int, optional
        Minimum number of connected nodes to consider for unwrapping. If masking
        separates the input data into disconnected sets of pixels, a source is
        selected for each connected set, provided that the number of nodes in
        the set is greater than n_conn_node_min. Must be nonnegative.
        (default: 0)
    n_major_prune : int, optional
        Number of major iterations between tree pruning operations. A smaller
        number causes pruning to occur more frequently. (default: 2000000000)
    prune_cost_thresh : int, optional
        Cost threshold for pruning the tree. A lower threshold prunes more
        aggressively. (default: 2000000000)
    """

    max_flow_inc: int = 4
    init_max_flow: int = 9999
    arc_max_flow_const: int = 3
    threshold: float = 0.001
    max_cost: float = 1000.0
    cost_scale: float = 100.0
    n_cycle: int = 200
    max_new_node_const: float = 0.0008
    max_n_flow_cycles: Optional[float] = None
    max_cycle_frac: float = 0.00001
    n_conn_node_min: int = 0
    n_major_prune: int = 2_000_000_000
    prune_cost_thresh: int = 2_000_000_000

    def tostring(self):
        """Convert to string in SNAPHU config file format."""
        s = ""
        s += f"MAXFLOW {self.max_flow_inc}\n"
        s += f"INITMAXFLOW {self.init_max_flow}\n"
        s += f"ARCMAXFLOWCONST {self.arc_max_flow_const}\n"
        s += f"THRESHOLD {self.threshold}\n"
        s += f"MAXCOST {self.max_cost}\n"
        s += f"COSTSCALE {self.cost_scale}\n"
        s += f"NSHORTCYCLE {self.n_cycle}\n"
        s += f"MAXNEWNODECONST {self.max_new_node_const}\n"
        if self.max_n_flow_cycles is not None:
            s += f"MAXNFLOWCYCLES {self.max_n_flow_cycles}\n"
        else:
            s += f"MAXCYCLEFRACTION {self.max_cycle_frac}\n"
        s += f"NCONNNODEMIN {self.n_conn_node_min}\n"
        s += f"NMAJORPRUNE {self.n_major_prune}\n"
        s += f"PRUNECOSTTHRESH {self.prune_cost_thresh}\n"
        return s

@dataclass(frozen=True)
class ConnCompParams:
    """Configuration parameters affecting the generation of connected component
    labels.

    Attributes
    ----------
    min_frac_area : float, optional
        Minimum size of a single connected component, as a fraction of the total
        number of pixels in the tile. (default: 0.01)
    cost_thresh : int, optional
        Cost threshold for connected components. Higher threshold will give
        smaller connected components. (default: 300)
    max_ncomps : int, optional
        Maximum number of connected components per tile. (default: 32)
    """

    min_frac_area: float = 0.01
    cost_thresh: int = 300
    max_ncomps: int = 32

    def tostring(self):
        """Convert to string in SNAPHU config file format."""
        s = ""
        s += f"MINCONNCOMPFRAC {self.min_frac_area}\n"
        s += f"CONNCOMPTHRESH {self.cost_thresh}\n"
        s += f"MAXNCOMPS {self.max_ncomps}\n"
        return s

@dataclass(frozen=True)
class CorrBiasModelParams:
    r"""Model parameters for estimating bias in sample correlation magnitude
    expected for zero true correlation

    The multilooked correlation magnitude of the interferometric pair
    :math:`z_1` and :math:`z_2` is commonly estimated as

    .. math::

        \rho = \left| \frac{ \sum_{i=1}^{N}{z_{1i} z_{2i} ^*} }
            { \sqrt{ \sum_{i=1}^{N}{ \left| z_{1i} \right| ^2 } }
            \sqrt{ \sum_{i=1}^{N}{ \left| z_{2i} \right| ^2 } } }  \right|

    where :math:`N` is the number of statistically independent looks. SNAPHU
    uses the estimated correlation coefficient to infer statistics of the
    interferometric phase.

    This estimator is biased with respect to the expected true correlation,
    particularly at lower correlation values. In order to compensate for this,
    SNAPHU models the expected biased correlation measure, given that true
    interferometric correlation is zero, as

    .. math:: \rho_0 = \frac{c_1}{N} + c_2

    where :math:`N` is the number of effective looks used to estimate the
    correlation and :math:`c_1` & :math:`c_2` are the model coefficients. This
    approximately matches the curves of Touzi et al [1]_.

    Attributes
    ----------
    c1, c2 : float, optional
        Correlation bias model parameters.
    min_corr_factor : float, optional
        Factor applied to expected minimum measured (biased) correlation
        coefficient. Values smaller than the threshold min_corr_factor * rho0
        are assumed to come from zero statistical correlation because of
        estimator bias. rho0 is the expected biased correlation measure if the
        true correlation is zero. (default: 1.25)

    References
    ----------
    .. [1] R. Touzi, A. Lopes, J. Bruniquel, and P. W. Vachon, "Coherence
       estimation for SAR imagery," IEEE Trans. Geosci. Remote Sens. 37, 135-149
       (1999).
    """

    c1: float = 1.3
    c2: float = 0.14
    min_corr_factor: float = 1.25

    def tostring(self):
        """Convert to string in SNAPHU config file format."""
        s = ""
        s += f"RHOSCONST1 {self.c1}\n"
        s += f"RHOSCONST2 {self.c2}\n"

        # This parameter has different names depending on which cost mode was
        # selected -- "RHOMINFACTOR" for "topo" mode, "DEFOTHRESHFACTOR" for
        # "defo" and "smooth" mode -- but the semantics & effect are the same.
        # We redundantly define both params here since it's harmless and avoids
        # introducing a dependency on cost mode.
        s += f"RHOMINFACTOR {self.min_corr_factor}\n"
        s += f"DEFOTHRESHFACTOR {self.min_corr_factor}\n"

        return s

@dataclass(frozen=True)
class PhaseStddevModelParams:
    r"""Model parameters for approximating phase standard deviation from
    correlation magnitude

    Interferometric phase standard deviation is modelled as

    .. math:: \sigma_{\phi} = \rho ^{ c_1 + c_2 * \log nlooks + c_3 * nlooks }

    where :math:`\rho` is the sample correlation magnitude and :math:`nlooks` is
    the effective number of looks. :math:`c_1`, :math:`c_2`, and :math:`c_3` are
    the model coefficients. This approximately matches the curves of [1]_.

    Attributes
    ----------
    c1, c2, c3 : float, optional
        Interferometric phase standard deviation model parameters.
    sigsq_min : int
        Minimum value of phase variance after quantization to integer values.
        Must be greater than zero to prevent division by zero. (default: 1)

    References
    ----------
    .. [1] J. S. Lee, K. W. Hoppel, S. A. Mango, and A. R. Miller, "Intensity
       and phase statistics of multilook polarimetric and interferometric SAR
       imagery," IEEE Trans. Geosci. Remote Sens. 32, 1017-1028 (1994).
    """

    c1: float = 0.4
    c2: float = 0.35
    c3: float = 0.06
    sigsq_min: int = 1

    def tostring(self):
        """Convert to string in SNAPHU config file format."""
        s = ""
        s += f"CSTD1 {self.c1}\n"
        s += f"CSTD2 {self.c2}\n"
        s += f"CSTD3 {self.c3}\n"
        s += f"SIGSQSHORTMIN {self.sigsq_min}\n"
        return s

@contextlib.contextmanager
def scratch_directory(d: Optional[os.PathLike] = None) -> pathlib.Path:
    """Context manager that creates a (possibly temporary) filesystem directory

    If the input is a path-like object, a directory will be created at the
    specified filesystem path if it did not already exist. The directory will
    persist after leaving the context manager scope.

    If the input is None, a temporary directory is created as though by
    `tempfile.TemporaryDirectory()`. Upon exiting the context manager scope, the
    directory and its contents are removed from the filesystem.

    Parameters
    ----------
    d : path-like or None, optional
        Scratch directory path. If None, a temporary directory is created.
        (default: None)

    Yields
    ------
    d : pathlib.Path
        Scratch directory path. If the input was None, the directory is removed
        from the filesystem upon exiting the context manager scope.
    """
    if d is None:
        try:
            d = tempfile.TemporaryDirectory()
            yield pathlib.Path(d.name)
        finally:
            d.cleanup()
    else:
        d = pathlib.Path(d)
        d.mkdir(parents=True, exist_ok=True)
        yield d

def to_flat_file(
    path: os.PathLike,
    raster: isce3.io.gdal.Raster,
    dtype: Optional[np.dtype] = None,
    batchsize: int = -1,
):
    """Write raster data to flat binary file.

    The output file is overwritten if it exists.

    Parameters
    ----------
    path : path-like
        Output filepath.
    raster : isce3.io.gdal.Raster
        Input raster.
    dtype : data-type or None, optional
        Output datatype. If None, use the input raster datatype. (default: None)
    batchsize : int, optional
        If this is a positive number, the data is copied serially in batches of
        this many rows to avoid holding the full dataset in memory at once.
        Otherwise, copy the full data array as a single batch. (default: -1)
    """
    if dtype is None:
        dtype = raster.data.dtype

    if batchsize < 1:
        batchsize = raster.length

    # Memory-map the output file.
    shape = (raster.length, raster.width)
    mmap = np.memmap(path, dtype=dtype, mode="w+", shape=shape)

    # Write data in batches.
    for i0 in range(0, raster.length, batchsize):
        i1 = i0 + batchsize
        mmap[i0:i1] = raster.data[i0:i1]

    # Explicitly flush to disk instead of waiting for the memory map to be
    # garbage-collected.
    mmap.flush()

def from_flat_file(
    path: os.PathLike,
    raster: isce3.io.gdal.Raster,
    dtype: Optional[np.dtype] = None,
    batchsize: int = -1,
):
    """Read raster data from flat binary file.

    Parameters
    ----------
    path : path-like
        Input filepath.
    raster : isce3.io.gdal.Raster
        Output raster.
    dtype : data-type or None, optional
        Input file datatype. If None, assume the same as the output raster
        datatype. (default: None)
    batchsize : int, optional
        If this is a positive number, the data is copied serially in batches of
        this many rows to avoid holding the full dataset in memory at once.
        Otherwise, copy the full data array as a single batch. (default: -1)
    """
    if dtype is None:
        dtype = raster.data.dtype

    if batchsize < 1:
        batchsize = raster.length

    # Memory-map the input file.
    shape = (raster.length, raster.width)
    mmap = np.memmap(path, dtype=dtype, mode="r", shape=shape)

    # Read data in batches.
    for i0 in range(0, raster.length, batchsize):
        i1 = i0 + batchsize
        raster.data[i0:i1] = mmap[i0:i1]

CostParams = Union[
    TopoCostParams, DefoCostParams, SmoothCostParams, PNormCostParams,
]
CostParams.__doc__ = """SNAPHU cost mode configuration parameters"""

def unwrap(
    unw: isce3.io.gdal.Raster,
    conncomp: isce3.io.gdal.Raster,
    igram: isce3.io.gdal.Raster,
    corr: isce3.io.gdal.Raster,
    nlooks: float,
    cost: str = "smooth",
    cost_params: Optional[CostParams] = None,
    init_method: str = "mcf",
    pwr: Optional[isce3.io.gdal.Raster] = None,
    mask: Optional[isce3.io.gdal.Raster] = None,
    unwest: Optional[isce3.io.gdal.Raster] = None,
    tiling_params: Optional[TilingParams] = None,
    solver_params: Optional[SolverParams] = None,
    conncomp_params: Optional[ConnCompParams] = None,
    corr_bias_model_params: Optional[CorrBiasModelParams] = None,
    phase_stddev_model_params: Optional[PhaseStddevModelParams] = None,
    scratchdir: Optional[os.PathLike] = None,
    debug: bool = False,
):
    r"""Performs 2-D phase unwrapping on an input interferogram using the SNAPHU
    algorithm.

    The algorithm attempts to estimate the unwrapped phase field by
    approximately solving a non-linear optimization problem using cost functions
    based on simple statistical models of the unwrapped phase gradients.

    The total cost is approximately minimized by applying a non-linear network
    flow solver based on the network simplex algorithm. An initial feasible
    solution is first computed before optimizing according to the specified cost
    mode.

    Different statistical cost functions may be applied depending on the
    application:
    - The "topo" cost mode generates cost functions for topographic SAR
    interferometry. The problem statistics are based on the assumption that the
    true unwrapped phase represents surface elevation. The input interferogram
    is assumed to be in radar (range-azimuth) coordinates for this mode.
    - The "defo" cost mode generates cost functions for deformation
    measurements. The problem statistics are based on the assumption that the
    true unwrapped phase represents surface displacement.
    - The "smooth" cost mode models the problem statistics based on the
    assumption that the true unwrapped phase represents a generic surface with
    no discontinuities.
    - The "p-norm" cost mode is not based on a statistical model but rather
    minimizes the :math:`L^p` norm of the difference between the unwrapped and
    wrapped phase gradients.

    The outputs include the unwrapped phase and a raster of connected component
    labels. Each connected component is a region of pixels in the solution that
    is believed to have been unwrapped in an internally self-consistent manner.
    Each distinct region is assigned a unique positive integer label. Pixels not
    belonging to any component are assigned a label of zero.

    The effective number of looks used to form the input correlation data must
    be provided in order to estimate interferometric phase statistics. The
    effective number of looks is an estimate of the number of statistically
    independent samples averaged in multilooked data, taking into account
    spatial correlation due to oversampling/filtering. It is approximately equal
    to

    .. math:: n_e = k_r k_a \frac{d_r d_a}{\rho_r \rho_a}

    where :math:`k_r` and :math:`k_a` are the number of looks in range and
    azimuth, :math:`d_r` and :math:`d_a` are the sample spacing in range and
    azimuth, and :math:`\rho_r` and :math:`\rho_a are the range and azimuth
    resolution.

    Unwrapping can be performed in tile mode to potentially speed up processing
    and make use of multiple processors in parallel, though this may result in
    processing artifacts at tile boundaries. The interferogram is partitioned
    into rectangular tiles, each of which is unwrapped independently before
    reassembly. The default behavior is to unwrap the full interferogram as a
    single tile.

    .. warning:: Currently, if tile mode is used and any connected component
    crosses spans multiple tiles, the assigned connected component label may be
    inconsistent across tiles.

    Parameters
    ----------
    unw : isce3.io.gdal.Raster
        Output raster for unwrapped phase, in radians. Must have the same
        dimensions as the input interferogram and GDT_Float32 datatype.
    conncomp : isce3.io.gdal.Raster
        Output connected component labels. Must have the same dimensions as the
        input interferogram and GDT_UInt32 datatype.
    igram : isce3.io.gdal.Raster
        Input interferogram. Must have GDT_CFloat32 datatype.
    corr : isce3.io.gdal.Raster
        Correlation magnitude, normalized to the interval [0, 1]. Must have the
        same dimensions as the input interferogram and GDT_Float32 datatype.
    nlooks : float
        Effective number of looks used to form the input correlation data.
    cost : {"topo", "defo", "smooth", "p-norm"}, optional
        Statistical cost mode. (default: "smooth")
    cost_params : CostParams or None, optional
        Configuration parameters for the specified cost mode. This argument is
        required for "topo" mode and optional for all other modes. If None, the
        default configuration parameters are used. (default: None)
    init_method: {"mst", "mcf"}, optional
        Algorithm used for initialization of unwrapped phase gradients.
        Supported algorithms include Minimum Spanning Tree ("mst") and Minimum
        Cost Flow ("mcf"). (default: "mcf")
    pwr : isce3.io.gdal.Raster or None, optional
        Average intensity of the two SLCs, in linear units (not dB). Only used
        in "topo" cost mode. If None, interferogram magnitude is used as
        intensity. Must have the same dimensions as the input interferogram and
        GDT_Float32 datatype. (default: None)
    mask : isce3.io.gdal.Raster or None, optional
        Binary mask of valid pixels. Zeros in this raster indicate interferogram
        pixels that should be masked out. Must have the same dimensions as the
        input interferogram and GDT_Byte datatype. (default: None)
    unwest : isce3.io.gdal.Raster or None, optional
        Initial estimate of unwrapped phase, in radians. This can be used to
        provide a coarse unwrapped estimate to guide the algorithm. Must have
        the same dimensions as the input interferogram and GDT_Float32 datatype.
        (default: None)
    tiling_params : TilingParams or None, optional
        Configuration parameters affecting scene tiling and parallel processing.
        If None, the default configuration parameters are used. (default: None)
    solver_params : SolverParams or None, optional
        Configuration parameters used by the network initialization and
        nonlinear network flow solver algorithms. If None, the default
        configuration parameters are used. (default: None)
    conncomp_params : ConnCompParams or None, optional
        Configuration parameters affecting the generation of connected component
        labels. If None, the default configuration parameters are used.
        (default: None)
    corr_bias_model_params : CorrBiasModelParams or None, optional
        Model parameters for estimating bias in sample correlation magnitude
        expected for zero true correlation. If None, the default model
        parameters are used. (default: None)
    phase_stddev_model_params : PhaseStddevModelParams or None, optional
        Model parameters for approximating phase standard deviation from
        correlation magnitude. If None, the default model parameters are used.
        (default: None)
    scratchdir : path-like or None, optional
        Scratch directory where intermediate processing artifacts are written.
        If the specified directory does not exist, it will be created. If None,
        a temporary directory will be created and automatically removed from the
        filesystem at the end of processing. Otherwise, the directory and its
        contents will not be cleaned up. (default: None)
    debug : bool, optional
        Dump intermediate data arrays to scratch directory for debugging?
        (default: False)

    See Also
    --------
    isce3.unwrap.ICU : Branch-cut-based phase unwrapping
    isce3.unwrap.Phass : Minimum Cost Flow-based phase unwrapping

    References
    ----------
    .. [1] C. W. Chen and H. A. Zebker, "Network approaches to two-dimensional
       phase unwrapping: intractability and two new algorithms," Journal of the
       Optical Society of America A, vol. 17, pp. 401-414 (2000).
    .. [2] C. W. Chen and H. A. Zebker, "Two-dimensional phase unwrapping with
       use of statistical models for cost functions in nonlinear optimization,"
       Journal of the Optical Society of America A, vol. 18, pp. 338-351 (2001).
    .. [3] C. W. Chen and H. A. Zebker, "Phase unwrapping for large SAR
       interferograms: Statistical segmentation and generalized network models,"
       IEEE Transactions on Geoscience and Remote Sensing, vol. 40, pp.
       1709-1719 (2002).
    """
    # Verify input & output raster datatypes.
    if unw.datatype != isce3.io.gdal.GDT_Float32:
        raise TypeError("unw raster must have GDT_Float32 datatype")
    if conncomp.datatype != isce3.io.gdal.GDT_UInt32:
        raise TypeError("conncomp raster must have GDT_UInt32 datatype")
    if igram.datatype != isce3.io.gdal.GDT_CFloat32:
        raise TypeError("igram raster must have GDT_CFloat32 datatype")
    if corr.datatype != isce3.io.gdal.GDT_Float32:
        raise TypeError("corr raster must have GDT_Float32 datatype")

    length, width = igram.length, igram.width

    # Check that raster dimensions are consistent.
    if (unw.length != length) or (unw.width != width):
        raise ValueError("unw raster dimensions must match interferogram")
    if (conncomp.length != length) or (conncomp.width != width):
        raise ValueError("conncomp raster dimensions must match interferogram")
    if (corr.length != length) or (corr.width != width):
        raise ValueError("corr raster dimensions must match interferogram")

    # Check specified number of effective looks.
    if nlooks < 1.0:
        raise ValueError("nlooks must be >= 1.0")

    # Generate a SNAPHU text configuration file to pass to the C++ code.
    configstr = ""
    configstr += f"LINELENGTH {width}\n"
    configstr += f"NCORRLOOKS {nlooks}\n"

    def cost_string():
        if cost == "topo":
            return "TOPO"
        if cost == "defo":
            return "DEFO"
        if cost == "smooth":
            return "SMOOTH"
        if cost == "p-norm":
            return "NOSTATCOSTS"
        raise ValueError(f"invalid cost mode '{cost}'")

    configstr += f"STATCOSTMODE {cost_string()}\n"

    def init_string():
        if init_method == "mst":
            return "MST"
        if init_method == "mcf":
            return "MCF"
        raise ValueError(f"invalid init method '{init_method}'")

    configstr += f"INITMETHOD {init_string()}\n"

    # Check cost mode-specific configuration params.
    if cost == "topo":
        # In "topo" mode, configuration params must be provided (there is no
        # default configuration).
        if not isinstance(cost_params, TopoCostParams):
            raise TypeError(
                "cost_params for 'topo' cost mode must be an "
                "instance of TopoCostParams"
            )
    elif cost == "defo":
        if cost_params is None:
            cost_params = DefoCostParams()
        if not isinstance(cost_params, DefoCostParams):
            raise TypeError("invalid cost_params for 'defo' cost mode")
    elif cost == "smooth":
        if cost_params is None:
            cost_params = SmoothCostParams()
        if not isinstance(cost_params, SmoothCostParams):
            raise TypeError("invalid cost_params for 'smooth' cost mode")
    elif cost == "p-norm":
        if cost_params is None:
            cost_params = PNormCostParams()
        if not isinstance(cost_params, PNormCostParams):
            raise TypeError("invalid cost_params for 'p-norm' cost mode")
    else:
        raise ValueError(f"invalid cost mode '{cost}'")

    configstr += cost_params.tostring()

    # Additional optional configuration parameters.
    if tiling_params is not None:
        configstr += tiling_params.tostring()
    if solver_params is not None:
        configstr += solver_params.tostring()
    if conncomp_params is not None:
        configstr += conncomp_params.tostring()

    # Curve-fitting coefficients.
    if corr_bias_model_params is not None:
        configstr += corr_bias_model_params.tostring()
    if phase_stddev_model_params is not None:
        configstr += phase_stddev_model_params.tostring()

    # Debug mode requires that a scratch directory is specified (otherwise, all
    # debug output would be automatically discarded anyway).
    if debug and (scratchdir is None):
        raise ValueError("scratchdir path must be specified if debug is True")

    # If no scratch directory was specified, make a temporary one. Otherwise,
    # create the directory if it doesn't already exist.
    with scratch_directory(scratchdir) as d:
        # SNAPHU expects flat binary data files, not GDAL rasters, as inputs &
        # outputs. Therefore, we create some intermediate files in the scratch
        # directory to pass to the backend code.

        # Output unwrapped data
        tmp_unw = d / "unw.f4"
        configstr += f"OUTFILE {tmp_unw.resolve()}\n"
        configstr += f"OUTFILEFORMAT FLOAT_DATA\n"

        # Output connected component labels
        tmp_conncomp = d / "conncomp.u4"
        configstr += f"CONNCOMPFILE {tmp_conncomp.resolve()}\n"
        configstr += f"CONNCOMPOUTTYPE UINT\n"

        # Input interferogram
        tmp_igram = d / "igram.c8"
        to_flat_file(tmp_igram, igram, batchsize=1024)
        configstr += f"INFILE {tmp_igram.resolve()}\n"
        configstr += f"INFILEFORMAT COMPLEX_DATA\n"

        # Input correlation magnitude
        tmp_corr = d / "corr.f4"
        to_flat_file(tmp_corr, corr, batchsize=1024)
        configstr += f"CORRFILE {tmp_corr.resolve()}\n"
        configstr += f"CORRFILEFORMAT FLOAT_DATA\n"

        # Input SLC intensity data
        if pwr is not None:
            if cost != "topo":
                raise ValueError("SLC intensity data is only used in 'topo' mode")
            if pwr.datatype != isce3.io.gdal.GDT_Float32:
                raise TypeError("pwr raster must have GDT_Float32 datatype")
            if (pwr.length != length) or (pwr.width != width):
                raise ValueError("pwr raster dimensions must match interferogram")

            tmp_pwr = d / "pwr.f4"
            to_flat_file(tmp_pwr, pwr, batchsize=1024)
            configstr += f"PWRFILE {tmp_pwr.resolve()}\n"
            configstr += f"AMPFILEFORMAT FLOAT_DATA\n"

        # Input data mask
        if mask is not None:
            if mask.datatype != isce3.io.gdal.GDT_Byte:
                raise TypeError("mask raster must have GDT_Byte datatype")
            if (mask.length != length) or (mask.width != width):
                raise ValueError("mask raster dimensions must match interferogram")

            tmp_mask = d / "mask.i1"
            to_flat_file(tmp_mask, mask, batchsize=1024)
            configstr += f"BYTEMASKFILE {tmp_mask.resolve()}\n"

        # Input unwrapped phase estimate
        if unwest is not None:
            if unwest.datatype != isce3.io.gdal.GDT_Float32:
                raise TypeError("unwest raster must have GDT_Float32 datatype")
            if (unwest.length != length) or (unwest.width != width):
                raise ValueError("unwest raster dimensions must match interferogram")

            tmp_unwest = d / "unwest.f4"
            to_flat_file(tmp_unwest, unwest, batchsize=1024)
            configstr += f"ESTIMATEFILE {tmp_unwest.resolve()}\n"
            configstr += f"ESTFILEFORMAT FLOAT_DATA\n"

        configstr += f"DEBUG {debug}\n"

        # Ensure that debug outputs are written to the scratch directory.
        if debug:
            configstr += f"INITFILE {(d / 'snaphu.init')}\n"
            configstr += f"FLOWFILE {(d / 'snaphu.flow')}\n"
            configstr += f"ROWCOSTFILE {(d / 'snaphu.rowcost')}\n"
            configstr += f"COLCOSTFILE {(d / 'snaphu.colcost')}\n"
            configstr += f"CORRDUMPFILE {(d / 'snaphu.corr')}\n"

        # Write config params to file.
        configpath = d / "snaphu.conf"
        configpath.write_text(configstr)

        # Run SNAPHU.
        _snaphu_unwrap(str(configpath))

        # Copy output data to GDAL rasters.
        from_flat_file(tmp_unw, unw, batchsize=1024)
        from_flat_file(tmp_conncomp, conncomp, batchsize=1024)


#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# 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.
#
# Author: Timothy Berkelbach <tim.berkelbach@gmail.com>
#

'''
Restricted QCISD implementation
The 4-index integrals are saved on disk entirely (without using any symmetry).

Note MO integrals are treated in chemist's notation

Ref:
'''

import numpy
import numpy as np

from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import rccsd_slow as rccsd
from pyscf.cc import rintermediates as imd
from pyscf import __config__

BLKMIN = getattr(__config__, 'cc_ccsd_blkmin', 4)
MEMORYMIN = getattr(__config__, 'cc_ccsd_memorymin', 2000)

def kernel(mycc, eris=None, t1=None, t2=None, max_cycle=50, tol=1e-8,
           tolnormt=1e-6, verbose=None):
    '''Same as ccsd.kernel with strings modified to correct the method name'''
    log = logger.new_logger(mycc, verbose)
    if eris is None:
        eris = mycc.ao2mo(mycc.mo_coeff)
    if t1 is None and t2 is None:
        t1, t2 = mycc.get_init_guess(eris)
    elif t2 is None:
        t2 = mycc.get_init_guess(eris)[1]

    cput1 = cput0 = (logger.process_clock(), logger.perf_counter())
    eold = 0
    eccsd = mycc.energy(t1, t2, eris)
    log.info('Init E_corr(QCISD) = %.15g', eccsd)

    if isinstance(mycc.diis, lib.diis.DIIS):
        adiis = mycc.diis
    elif mycc.diis:
        adiis = lib.diis.DIIS(mycc, mycc.diis_file, incore=mycc.incore_complete)
        adiis.space = mycc.diis_space
    else:
        adiis = None

    conv = False
    for istep in range(max_cycle):
        t1new, t2new = mycc.update_amps(t1, t2, eris)
        tmpvec = mycc.amplitudes_to_vector(t1new, t2new)
        tmpvec -= mycc.amplitudes_to_vector(t1, t2)
        normt = numpy.linalg.norm(tmpvec)
        tmpvec = None
        if mycc.iterative_damping < 1.0:
            alpha = mycc.iterative_damping
            t1new = (1-alpha) * t1 + alpha * t1new
            t2new *= alpha
            t2new += (1-alpha) * t2
        t1, t2 = t1new, t2new
        t1new = t2new = None
        t1, t2 = mycc.run_diis(t1, t2, istep, normt, eccsd-eold, adiis)
        eold, eccsd = eccsd, mycc.energy(t1, t2, eris)
        log.info('cycle = %d  E_corr(QCISD) = %.15g  dE = %.9g  norm(t1,t2) = %.6g',
                 istep+1, eccsd, eccsd - eold, normt)
        cput1 = log.timer('QCISD iter', *cput1)
        if abs(eccsd-eold) < tol and normt < tolnormt:
            conv = True
            break
    log.timer('QCISD', *cput0)
    return conv, eccsd, t1, t2

def update_amps(cc, t1, t2, eris):
    # Ref: Hirata et al., J. Chem. Phys. 120, 2581 (2004) Eqs.(35)-(36)
    nocc, nvir = t1.shape
    fock = eris.fock

    fov = fock[:nocc,nocc:].copy()
    foo = fock[:nocc,:nocc].copy()
    fvv = fock[nocc:,nocc:].copy()

    Foo = imd.cc_Foo(0*t1,t2,eris)
    Fvv = imd.cc_Fvv(0*t1,t2,eris)
    Fov = imd.cc_Fov(t1,t2,eris)

    Foo -= np.diag(np.diag(foo))
    Fvv -= np.diag(np.diag(fvv))

    # T1 equation
    t1new = np.asarray(fov).conj().copy()
    t1new +=   lib.einsum('ac,ic->ia', Fvv, t1)
    t1new +=  -lib.einsum('ki,ka->ia', Foo, t1)
    t1new += 2*lib.einsum('kc,kica->ia', Fov, t2)
    t1new +=  -lib.einsum('kc,ikca->ia', Fov, t2)
    t1new += 2*lib.einsum('kcai,kc->ia', eris.ovvo, t1)
    t1new +=  -lib.einsum('kiac,kc->ia', eris.oovv, t1)
    eris_ovvv = np.asarray(eris.ovvv)
    t1new += 2*lib.einsum('kdac,ikcd->ia', eris_ovvv, t2)
    t1new +=  -lib.einsum('kcad,ikcd->ia', eris_ovvv, t2)
    t1new +=-2*lib.einsum('kilc,klac->ia', eris.ooov, t2)
    t1new +=   lib.einsum('likc,klac->ia', eris.ooov, t2)

    # T2 equation
    t2new = np.asarray(eris.ovov).conj().transpose(0,2,1,3).copy()
    Loo = imd.Loo(0*t1, t2, eris)
    Lvv = imd.Lvv(0*t1, t2, eris)
    Loo -= np.diag(np.diag(foo))
    Lvv -= np.diag(np.diag(fvv))
    Woooo = imd.cc_Woooo(0*t1, t2, eris)
    Wvoov = imd.cc_Wvoov(0*t1, t2, eris)
    Wvovo = imd.cc_Wvovo(0*t1, t2, eris)
    Wvvvv = imd.cc_Wvvvv(0*t1, t2, eris)
    t2new += lib.einsum('klij,klab->ijab', Woooo, t2)
    t2new += lib.einsum('abcd,ijcd->ijab', Wvvvv, t2)
    tmp = lib.einsum('ac,ijcb->ijab', Lvv, t2)
    t2new += (tmp + tmp.transpose(1,0,3,2))
    tmp = lib.einsum('ki,kjab->ijab', Loo, t2)
    t2new -= (tmp + tmp.transpose(1,0,3,2))
    tmp  = 2*lib.einsum('akic,kjcb->ijab', Wvoov, t2)
    tmp -=   lib.einsum('akci,kjcb->ijab', Wvovo, t2)
    t2new += (tmp + tmp.transpose(1,0,3,2))
    tmp = lib.einsum('akic,kjbc->ijab', Wvoov, t2)
    t2new -= (tmp + tmp.transpose(1,0,3,2))
    tmp = lib.einsum('bkci,kjac->ijab', Wvovo, t2)
    t2new -= (tmp + tmp.transpose(1,0,3,2))

    tmp2 = np.asarray(eris.ovvv).conj().transpose(1,3,0,2)
    tmp = lib.einsum('abic,jc->ijab', tmp2, t1)
    t2new += (tmp + tmp.transpose(1,0,3,2))
    tmp2 = np.asarray(eris.ooov).transpose(3,1,2,0).conj()
    tmp = lib.einsum('akij,kb->ijab', tmp2, t1)
    t2new -= (tmp + tmp.transpose(1,0,3,2))

    mo_e = eris.fock.diagonal().real
    eia = mo_e[:nocc,None] - mo_e[None,nocc:]
    eijab = lib.direct_sum('ia,jb->ijab',eia,eia)
    t1new /= eia
    t2new /= eijab

    return t1new, t2new

class QCISD(rccsd.RCCSD):
    '''restricted QCISD
    '''

    def kernel(self, t1=None, t2=None, eris=None):
        return self.qcisd(t1, t2, eris)
    def qcisd(self, t1=None, t2=None, eris=None):
        assert(self.mo_coeff is not None)
        assert(self.mo_occ is not None)

        if self.verbose >= logger.WARN:
            self.check_sanity()
        self.dump_flags()

        if eris is None:
            eris = self.ao2mo(self.mo_coeff)

        self.e_hf = getattr(eris, 'e_hf', None)
        if self.e_hf is None:
            self.e_hf = self._scf.e_tot

        self.converged, self.e_corr, self.t1, self.t2 = \
                kernel(self, eris, t1, t2, max_cycle=self.max_cycle,
                       tol=self.conv_tol, tolnormt=self.conv_tol_normt,
                       verbose=self.verbose)
        self._finalize()
        return self.e_corr, self.t1, self.t2

    def energy(self, t1=None, t2=None, eris=None):
        return rccsd.energy(self, t1*0, t2, eris)

    update_amps = update_amps

    def qcisd_t(self, t1=None, t2=None, eris=None):
        from pyscf.cc import qcisd_t_slow as qcisd_t
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if eris is None: eris = self.ao2mo(self.mo_coeff)
        return qcisd_t.kernel(self, eris, t1, t2, self.verbose)

    def density_fit(self, auxbasis=None, with_df=None):
        raise NotImplementedError

if __name__ == '__main__':
    from pyscf import gto, scf

    mol = gto.Mole()
    mol.atom = """C  0.000  0.000  0.000
                  H  0.637  0.637  0.637
                  H -0.637 -0.637  0.637
                  H -0.637  0.637 -0.637
                  H  0.637 -0.637 -0.637"""
    mol.basis = 'cc-pvdz'
    mol.verbose = 7
    mol.spin = 0
    mol.build()
    mf = scf.RHF(mol).run(conv_tol=1e-14)

    mycc = QCISD(mf, frozen=1)
    ecc, t1, t2 = mycc.kernel()
    print(mycc.e_tot - -40.383989)
    et = mycc.qcisd_t()
    print(mycc.e_tot+et - -40.387679)


"""
Energy-conserving solver for the binary accretion problem in 2D.
"""

from typing import NamedTuple
from logging import getLogger
from sailfish.kernel.library import Library
from sailfish.kernel.system import get_array_module, execution_context, num_devices
from sailfish.mesh import PlanarCartesian2DMesh
from sailfish.physics.circumbinary import Physics, EquationOfState, ViscosityModel
from sailfish.solver import SolverBase
from sailfish.subdivide import subdivide, concat_on_host, lazy_reduce

logger = getLogger(__name__)

class Options(NamedTuple):
    pressure_floor: float = 1e-12
    density_floor: float = 1e-10
    velocity_ceiling: float = 1e16
    mach_ceiling: float = 1e5

def initial_condition(setup, mesh, time):
    """
    Generate a 2D array of primitive data from a mesh and a setup.
    """
    import numpy as np

    ni, nj = mesh.shape
    primitive = np.zeros([ni, nj, 4])

    for i in range(ni):
        for j in range(nj):
            setup.primitive(time, mesh.cell_coordinates(i, j), primitive[i, j])

    return primitive

class Patch:
    """
    Holds the array buffer state for the solution on a subset of the
    solution domain.
    """

    def __init__(
        self,
        time,
        primitive,
        mesh,
        index_range,
        physics,
        options,
        buffer_outer_radius,
        buffer_surface_density,
        buffer_surface_pressure,
        lib,
        xp,
        execution_context,
    ):
        i0, i1 = index_range
        ni, nj = i1 - i0, mesh.shape[1]
        self.lib = lib
        self.mesh = mesh
        self.xp = xp
        self.execution_context = execution_context
        self.time = self.time0 = time
        self.shape = (i1 - i0, nj)  # not including guard zones
        self.physics = physics
        self.options = options
        self.xl, self.yl = mesh.vertex_coordinates(i0, 0)
        self.xr, self.yr = mesh.vertex_coordinates(i1, nj)
        self.buffer_outer_radius = buffer_outer_radius
        self.buffer_surface_density = buffer_surface_density
        self.buffer_surface_pressure = buffer_surface_pressure

        with self.execution_context:
            self.wavespeeds = self.xp.zeros(primitive.shape[:2])
            self.primitive1 = self.xp.array(primitive)
            self.primitive2 = self.xp.array(primitive)
            self.conserved0 = self.xp.zeros(primitive.shape)

    def point_mass_source_term(self, which_mass):
        if which_mass not in (1, 2):
            raise ValueError("the mass must be either 1 or 2")

        m1, m2 = self.physics.point_masses(self.time)
        with self.execution_context:
            cons_rate = self.xp.zeros_like(self.conserved0)
            self.lib.cbdgam_2d_point_mass_source_term[self.shape](
                self.xl,
                self.xr,
                self.yl,
                self.yr,
                m1.position_x,
                m1.position_y,
                m1.velocity_x,
                m1.velocity_y,
                m1.mass,
                m1.softening_length,
                m1.sink_rate,
                m1.sink_radius,
                m1.sink_model.value,
                m2.position_x,
                m2.position_y,
                m2.velocity_x,
                m2.velocity_y,
                m2.mass,
                m2.softening_length,
                m2.sink_rate,
                m2.sink_radius,
                m2.sink_model.value,
                which_mass,
                self.primitive1,
                cons_rate,
                int(self.physics.constant_softening),
                self.physics.gamma_law_index,
            )
            return cons_rate

    def maximum_wavespeed(self):
        with self.execution_context:
            self.lib.cbdgam_2d_wavespeed[self.shape](
                self.primitive1,
                self.wavespeeds,
                self.physics.gamma_law_index,
            )
            return self.wavespeeds.max()

    def recompute_conserved(self):
        with self.execution_context:
            return self.lib.cbdgam_2d_primitive_to_conserved[self.shape](
                self.primitive1,
                self.conserved0,
                self.physics.gamma_law_index,
            )

    def advance_rk(self, rk_param, dt):
        m1, m2 = self.physics.point_masses(self.time)
        buffer_central_mass = m1.mass + m2.mass
        buffer_surface_density = self.buffer_surface_density
        buffer_surface_pressure = self.buffer_surface_pressure

        with self.execution_context:
            self.lib.cbdgam_2d_advance_rk[self.shape](
                self.xl,
                self.xr,
                self.yl,
                self.yr,
                self.conserved0,
                self.primitive1,
                self.primitive2,
                self.physics.gamma_law_index,
                buffer_surface_density,
                buffer_surface_pressure,
                buffer_central_mass,
                self.physics.buffer_driving_rate,
                self.buffer_outer_radius,
                self.physics.buffer_onset_width,
                int(self.physics.buffer_is_enabled),
                m1.position_x,
                m1.position_y,
                m1.velocity_x,
                m1.velocity_y,
                m1.mass,
                m1.softening_length,
                m1.sink_rate,
                m1.sink_radius,
                m1.sink_model.value,
                m2.position_x,
                m2.position_y,
                m2.velocity_x,
                m2.velocity_y,
                m2.mass,
                m2.softening_length,
                m2.sink_rate,
                m2.sink_radius,
                m2.sink_model.value,
                self.physics.alpha,
                rk_param,
                dt,
                self.options.velocity_ceiling,
                self.physics.cooling_coefficient,
                self.options.mach_ceiling,
                self.options.density_floor,
                self.options.pressure_floor,
                int(self.physics.constant_softening),
            )

        self.time = self.time0 * rk_param + (self.time + dt) * (1.0 - rk_param)
        self.primitive1, self.primitive2 = self.primitive2, self.primitive1

    def new_iteration(self):
        self.time0 = self.time
        self.recompute_conserved()

    @property
    def primitive(self):
        return self.primitive1

class Solver(SolverBase):
    """
    Adapter class to drive the cbdgam_2d C extension module.
    """

    def __init__(
        self,
        setup=None,
        mesh=None,
        time=0.0,
        solution=None,
        num_patches=1,
        mode="cpu",
        physics=dict(),
        options=dict(),
    ):
        import numpy as np

        self._physics = physics = Physics(**physics)
        self._options = options = Options(**options)

        if type(mesh) is not PlanarCartesian2DMesh:
            raise ValueError("solver only supports 2D cartesian mesh")

        if setup.boundary_condition != "outflow":
            raise ValueError("solver only supports outflow boundary condition")

        if physics.viscosity_model not in (
            ViscosityModel.NONE,
            ViscosityModel.CONSTANT_ALPHA,
        ):
            raise ValueError("solver only supports constant-nu viscosity")

        if physics.eos_type != EquationOfState.GAMMA_LAW:
            raise ValueError("solver only supports isothermal equation of states")

        xp = get_array_module(mode)
        ng = 2  # number of guard zones
        nq = 4  # number of conserved quantities
        with open(__file__.replace(".py", ".c")) as f:
            code = f.read()
        lib = Library(code, mode=mode, debug=True)

        logger.info(f"initiate with time={time:0.4f}")
        logger.info(f"subdivide grid over {num_patches} patches")
        logger.info(f"mesh is {mesh}")
        logger.info(f"boundary condition is outflow")

        self.mesh = mesh
        self.setup = setup
        self.num_guard = ng
        self.num_cons = nq
        self.xp = xp
        self.patches = []
        ni, nj = mesh.shape
        self.domain_radius = self.mesh.x1
        self.buffer_onset_width = 0.1

        if solution is None:
            primitive = initial_condition(setup, mesh, time)
        else:
            primitive = solution

        if physics.buffer_is_enabled:
            # Here we sample the initial condition at the buffer onset radius
            # to determine the disk surface density at the radius where the
            # buffer begins to ramp up. This procedure makes sense as long as
            # the initial condition is axisymmetric.
            buffer_prim = [0.0] * 4
            buffer_outer_radius = mesh.x1  # this assumes the mesh is a centered squared
            buffer_onset_radius = buffer_outer_radius - physics.buffer_onset_width
            setup.primitive(time, [buffer_onset_radius, 0.0], buffer_prim)
            buffer_surface_density = buffer_prim[0]
            buffer_surface_pressure = buffer_prim[3]
        else:
            buffer_outer_radius = 0.0
            buffer_surface_density = 0.0
            buffer_surface_pressure = 0.0

        for n, (a, b) in enumerate(subdivide(ni, num_patches)):
            prim = np.zeros([b - a + 2 * ng, nj + 2 * ng, nq])
            prim[ng:-ng, ng:-ng] = primitive[a:b]
            patch = Patch(
                time,
                prim,
                mesh,
                (a, b),
                physics,
                options,
                buffer_outer_radius,
                buffer_surface_density,
                buffer_surface_pressure,
                lib,
                xp,
                execution_context(mode, device_id=n % num_devices(mode)),
            )
            self.patches.append(patch)

    @property
    def solution(self):
        return self.primitive

    @property
    def primitive(self):
        return concat_on_host(
            [p.primitive for p in self.patches], (self.num_guard, self.num_guard)
        )

    @property
    def time(self):
        return self.patches[0].time

    @property
    def options(self):
        return self._options._asdict()

    @property
    def physics(self):
        return self._physics._asdict()

    @property
    def recommended_cfl(self):
        return 0.1

    @property
    def maximum_cfl(self):
        return 0.4

    def maximum_wavespeed(self):
        return lazy_reduce(
            max,
            float,
            (patch.maximum_wavespeed for patch in self.patches),
            (patch.execution_context for patch in self.patches),
        )

    def advance(self, dt):
        self.new_iteration()
        self.advance_rk(0.0, dt)
        self.advance_rk(0.5, dt)

    def advance_rk(self, rk_param, dt):
        self.set_bc("primitive1")
        for patch in self.patches:
            patch.advance_rk(rk_param, dt)

    def set_bc(self, array):
        ng = self.num_guard
        num_patches = len(self.patches)
        for i0 in range(num_patches):
            il = (i0 + num_patches - 1) % num_patches
            ir = (i0 + num_patches + 1) % num_patches
            pl = getattr(self.patches[il], array)
            pc = getattr(self.patches[i0], array)
            pr = getattr(self.patches[ir], array)
            self.set_bc_patch(pl, pc, pr, i0)

    def set_bc_patch(self, pl, pc, pr, patch_index):
        ni, nj = self.mesh.shape
        ng = self.num_guard

        # 1. write to the guard zones of pc, the internal BC
        pc[:+ng] = pl[-2 * ng : -ng]
        pc[-ng:] = pr[+ng : +2 * ng]

        # 2. Set outflow BC on the left/right patch edges
        if patch_index == 0:
            for i in range(ng):
                pc[i] = pc[ng]
        if patch_index == len(self.patches) - 1:
            for i in range(pc.shape[0] - ng, pc.shape[0]):
                pc[i] = pc[-ng - 1]

        # 3. Set outflow BC on bottom and top edges
        for i in range(ng):
            pc[:, i] = pc[:, ng]

        for i in range(pc.shape[1] - ng, pc.shape[1]):
            pc[:, i] = pc[:, -ng - 1]

    def new_iteration(self):
        for patch in self.patches:
            patch.new_iteration()


"""Wrapper for performing an electronic circular dichroism (ECD)
calculation."""

import numpy as np

from pyresponse.constants import HARTREE_TO_EV, HARTREE_TO_INVCM, alpha, esuecd
from pyresponse.core import Program
from pyresponse.molecular_property import TransitionProperty
from pyresponse.operators import Operator
from pyresponse.td import TDHF
from pyresponse.utils import form_indices_zero

class ECD(TransitionProperty):
    """Wrapper for performing an electronic circular dichroism (ECD)
    calculation.
    """

    def __init__(
        self,
        program: Program,
        program_obj,
        driver: TDHF,
        mocoeffs: np.ndarray,
        moenergies: np.ndarray,
        occupations: np.ndarray,
        *,
        do_tda: bool = False,
        do_dipvel: bool = False,
    ) -> None:
        super().__init__(
            program, program_obj, driver, mocoeffs, moenergies, occupations, do_tda=do_tda
        )
        self.do_dipvel = do_dipvel

    def form_operators(self) -> None:

        if self.program == Program.PySCF:
            from pyresponse.pyscf import integrals

            integral_generator = integrals.IntegralsPyscf(self.program_obj)
        elif self.program == Program.Psi4:
            from pyresponse.psi4 import integrals

            integral_generator = integrals.IntegralsPsi4(self.program_obj)
        else:
            raise RuntimeError

        operator_angmom = Operator(
            label="angmom", is_imaginary=True, is_spin_dependent=False, triplet=False
        )
        operator_angmom.ao_integrals = integral_generator.integrals(integrals.ANGMOM_COMMON_GAUGE)
        self.driver.add_operator(operator_angmom)

        operator_diplen = Operator(
            label="dipole", is_imaginary=False, is_spin_dependent=False, triplet=False
        )
        operator_diplen.ao_integrals = integral_generator.integrals(integrals.DIPOLE)
        self.driver.add_operator(operator_diplen)

        if self.do_dipvel:
            operator_dipvel = Operator(
                label="dipvel", is_imaginary=True, is_spin_dependent=False, triplet=False
            )
            operator_dipvel.ao_integrals = integral_generator.integrals(integrals.DIPVEL)
            self.driver.add_operator(operator_dipvel)

    def form_results(self) -> None:

        operator_angmom = self.driver.solver.operators[0]
        operator_diplen = self.driver.solver.operators[1]
        assert len(operator_angmom.transition_moments) == len(operator_diplen.transition_moments)
        nstates = len(operator_diplen.transition_moments)
        rotational_strengths_diplen = []
        rotational_strengths_dipvel = []
        if self.do_dipvel:
            assert len(self.driver.solver.operators) == 3
            operator_dipvel = self.driver.solver.operators[2]
            assert len(operator_dipvel.transition_moments) == nstates
        for stateidx in range(nstates):
            print("-" * 78)
            eigval = self.driver.solver.eigvals[stateidx].real
            rotstr_diplen = (
                esuecd
                * (-1 / 2)
                * np.dot(
                    operator_diplen.transition_moments[stateidx],
                    operator_angmom.transition_moments[stateidx],
                )
            )
            print("length  ", rotstr_diplen)
            rotational_strengths_diplen.append(rotstr_diplen)
            if self.do_dipvel:
                rotstr_dipvel = (
                    esuecd
                    * (-1 / 2)
                    * np.dot(
                        operator_dipvel.transition_moments[stateidx] / eigval,
                        operator_angmom.transition_moments[stateidx],
                    )
                )
                print("velocity", rotstr_dipvel)
                rotational_strengths_dipvel.append(rotstr_dipvel)
        self.rotational_strengths_diplen = np.array(rotational_strengths_diplen)
        if self.do_dipvel:
            self.rotational_strengths_dipvel = np.array(rotational_strengths_dipvel)

    def print_results_nwchem(self) -> str:
        excitation_block = self.driver.print_results_nwchem()
        lines = [excitation_block]
        energies = self.driver.solver.eigvals.real
        energies_ev = energies * HARTREE_TO_EV
        op_diplen = self.driver.solver.operators[1]
        tmom_diplen = op_diplen.transition_moments
        etoscslen = op_diplen.total_oscillator_strengths
        op_angmom = self.driver.solver.operators[0]
        tmom_angmom = op_angmom.transition_moments
        rotstrlen = self.rotational_strengths_diplen
        if self.do_dipvel:
            op_dipvel = self.driver.solver.operators[2]
            rotstrvel = self.rotational_strengths_dipvel
            tmom_dipvel = op_dipvel.transition_moments
            etoscsvel = op_dipvel.total_oscillator_strengths
        nstates = len(energies)
        for state in range(nstates):
            lines.append(
                "  ----------------------------------------------------------------------------"
            )
            lines.append(
                f"  Root {state + 1:>3d} singlet a{energies[state]:>25.9f} a.u.{energies_ev[state]:>22.4f} eV"
            )
            lines.append(
                "  ----------------------------------------------------------------------------"
            )
            lines.append(
                f"     Transition Moments    X{tmom_diplen[state, 0]:>9.5f}   Y{tmom_diplen[state, 1]:>9.5f}   Z{tmom_diplen[state, 2]:>9.5f}"
            )
            ## TODO these require second moment (length) integrals
            ## lines.append(f'     Transition Moments   XX -0.28379  XY  0.08824  XZ -0.17416')
            ## lines.append(f'     Transition Moments   YY -0.40247  YZ -0.45981  ZZ  0.59211')
            lines.append(f"     Dipole Oscillator Strength {etoscslen[state]:>31.5f}")
            lines.append("")
            lines.append("     Electric Transition Dipole:")
            lines.append(
                f"            X{tmom_diplen[state, 0]:>13.7f}   Y{tmom_diplen[state, 1]:>13.7f}   Z{tmom_diplen[state, 2]:>13.7f}"
            )
            lines.append("     Magnetic Transition Dipole (Length):")
            lines.append(
                f"            X{tmom_angmom[state, 0]:>13.7f}   Y{tmom_angmom[state, 1]:>13.7f}   Z{tmom_angmom[state, 2]:>13.7f}"
            )
            lines.append("     Magnetic Transition Dipole * 1/c :")
            lines.append(
                f"            X{tmom_angmom[state, 0] * alpha:>13.7f}   Y{tmom_angmom[state, 1] * alpha:>13.7f}   Z{tmom_angmom[state, 2] * alpha:>13.7f}"
            )
            lines.append(f"     Rotatory Strength (1E-40 esu**2cm**2):{rotstrlen[state]:>21.7f}")
            lines.append("")
            if self.do_dipvel:
                lines.append("     Electric Transition Dipole (velocity representation):")
                lines.append(
                    f"            X{tmom_dipvel[state, 0]:>13.7f}   Y{tmom_dipvel[state, 1]:>13.7f}   Z{tmom_dipvel[state, 2]:>13.7f}"
                )
                # lines.append(f'     Oscillator Strength (velocity repr.) :            0.0069989')
                lines.append(
                    f"     Oscillator Strength (velocity repr.) :{etoscsvel[state]:>21.7f}"
                )
                # lines.append(f'     Oscillator Strength (mixed repr.   ) :            0.0074981')
                # lines.append(f'     Oscillator Strength (mixed repr.   ) :{>21.7f}')
                lines.append(
                    f"     Rotatory Strength   (velocity repr.) :{rotstrvel[state]:>21.7f}"
                )
                lines.append("")
            # lines.append(str(self.driver.solver.eigvecs[:, state]))
            # lines.append(str(self.driver.solver.eigvecs_normed[:, state]))

        return "\n".join(lines)

    def print_results_orca(self) -> str:
        excitation_block = self.driver.print_results_orca()
        lines = [excitation_block]
        energies = self.driver.solver.eigvals.real
        energies_to_invcm = energies * HARTREE_TO_INVCM
        energies_to_nm = 10000000 / energies_to_invcm
        op_diplen = self.driver.solver.operators[1]
        etoscslen = op_diplen.total_oscillator_strengths
        tmom_diplen = op_diplen.transition_moments
        t2_diplen = np.asarray(
            [np.dot(tmom_diplen[x], tmom_diplen[x]) for x in range(len(tmom_diplen))]
        )
        rotstrlen = self.rotational_strengths_diplen
        op_angmom = self.driver.solver.operators[0]
        tmom_angmom = op_angmom.transition_moments
        if self.do_dipvel:
            op_dipvel = self.driver.solver.operators[2]
            rotstrvel = self.rotational_strengths_dipvel
            etoscsvel = op_dipvel.total_oscillator_strengths
            tmom_dipvel = op_dipvel.transition_moments
            t2_dipvel = np.asarray(
                [np.dot(tmom_dipvel[x], tmom_dipvel[x]) for x in range(len(tmom_dipvel))]
            )
        nstates = len(energies)
        lines.append(
            "-----------------------------------------------------------------------------"
        )
        lines.append("         ABSORPTION SPECTRUM VIA TRANSITION ELECTRIC DIPOLE MOMENTS")
        lines.append(
            "-----------------------------------------------------------------------------"
        )
        lines.append("State   Energy  Wavelength   fosc         T2         TX        TY        TZ")
        lines.append("        (cm-1)    (nm)                  (au**2)     (au)      (au)      (au)")
        lines.append(
            "-----------------------------------------------------------------------------"
        )
        for state in range(nstates):
            lines.append(
                f"{state + 1:>4d}{energies_to_invcm[state]:>10.1f}{energies_to_nm[state]:>9.1f}{etoscslen[state]:>14.9f}{t2_diplen[state]:>10.5f}{tmom_diplen[state, 0]:>10.5f}{tmom_diplen[state, 1]:>10.5f}{tmom_diplen[state, 2]:>10.5f}"
            )
        lines.append("")
        lines.append(
            "-----------------------------------------------------------------------------"
        )
        lines.append("         ABSORPTION SPECTRUM VIA TRANSITION VELOCITY DIPOLE MOMENTS")
        lines.append(
            "-----------------------------------------------------------------------------"
        )
        lines.append("State   Energy  Wavelength   fosc         P2         PX        PY        PZ")
        lines.append("        (cm-1)    (nm)                  (au**2)     (au)      (au)      (au)")
        lines.append(
            "-----------------------------------------------------------------------------"
        )
        for state in range(nstates):
            lines.append(
                f"{state + 1:>4d}{energies_to_invcm[state]:>10.1f}{energies_to_nm[state]:>9.1f}{etoscsvel[state]:>14.9f}{t2_dipvel[state]:>10.5f}{tmom_dipvel[state, 0]:>10.5f}{tmom_dipvel[state, 1]:>10.5f}{tmom_dipvel[state, 2]:>10.5f}"
            )
        lines.append("")
        lines.append("-------------------------------------------------------------------")
        lines.append("                             CD SPECTRUM")
        lines.append("-------------------------------------------------------------------")
        lines.append("State   Energy Wavelength       R         MX        MY        MZ")
        lines.append("        (cm-1)   (nm)       (1e40*cgs)   (au)      (au)      (au)")
        lines.append("-------------------------------------------------------------------")
        for state in range(nstates):
            lines.append(
                f"{state + 1:>4d}{energies_to_invcm[state]:>10.1f}{energies_to_nm[state]:>9.1f}{rotstrlen[state]:>13.5f}{tmom_angmom[state, 0]:>10.5f}{tmom_angmom[state, 1]:>10.5f}{tmom_angmom[state, 2]:>10.5f}"
            )
        lines.append("")
        return "\n".join(lines)

    _HAMILTONIAN_PREFIX_QCHEM = {"tda": "", "rpa": "X: "}

    # TODO cutoff taken from ORCA, check the source code to see the
    # real criterion
    def print_results_qchem(self, cutoff: float = 0.01) -> str:
        energies = self.driver.solver.eigvals.real
        energies_ev = energies * HARTREE_TO_EV
        op_diplen = self.driver.solver.operators[1]
        tmom_diplen = op_diplen.transition_moments
        etoscslen = op_diplen.total_oscillator_strengths
        nocc_tot, nvirt_tot, _, _ = self.driver.solver.occupations
        indices = form_indices_zero(nocc_tot, nvirt_tot)
        eigvecs = self.driver.solver.eigvecs
        square_eigvecs = np.power(eigvecs, 2)
        lines = []
        lines.append(" ---------------------------------------------------")
        lines.append(
            f"               {self.driver._HAMILTONIAN_MAP_ORCA[self.driver.hamiltonian]} Excitation Energies              "
        )
        lines.append(" ---------------------------------------------------")
        lines.append("")
        nstates = len(energies)
        for state in range(nstates):
            lines.append(
                f" Excited state{state + 1:>4d}: excitation energy (eV) ={energies_ev[state]:>10.4f}"
            )
            lines.append(f" Total energy for state{state + 1:>3d}:{0:>31.8f} au")
            lines.append(f"    Multiplicity: {self.driver._SPIN_MAP_QCHEM[self.driver.spin]}")
            lines.append(
                f"    Trans. Mom.:{tmom_diplen[state, 0]:>8.4f} X{tmom_diplen[state, 1]:>9.4f} Y{tmom_diplen[state, 2]:>9.4f} Z"
            )
            lines.append(f"    Strength   :{etoscslen[state]:>17.10f}")
            eigvec_state = eigvecs[:, state]
            square_eigvec_state = square_eigvecs[:, state]
            mask = square_eigvec_state > cutoff
            coeffs_cutoff = eigvec_state[mask]
            mask_indices = np.array([p for (p, b) in enumerate(mask) if b])
            for i in range(len(coeffs_cutoff)):
                iocc, ivirt = indices[mask_indices[i]]
                lines.append(
                    f"    {self._HAMILTONIAN_PREFIX_QCHEM[self.driver.hamiltonian]}D({iocc + 1:>3d}) --> V({ivirt + 1:>3d}) amplitude ={coeffs_cutoff[i]:>8.4f}"
                )
            lines.append("")
        return "\n".join(lines)