"""
adaptive.py

functions relating to simulating an AO system with Proper
mostly code copied from the original MEDIS from Rupert

Generally, the optical prescription will call the deformable_mirror function, which will compile all information and
run sequentially all functions related to creating the adaptive optic correction (main AO functionality, relating to
atmospheric and common-path aberrations) as well as using or not CDI probes, DM corrections or errors, etc

TODO
    Add astrogrid pattern functionality from MEDIS0

"""

import numpy as np
from scipy import interpolate, ndimage
from inspect import getframeinfo, stack
from skimage.restoration import unwrap_phase
import matplotlib.pylab as plt
import proper
import datetime

from medis.params import sp, tp, ap
from medis.CDI import cdi, config_probe
from medis.optics import check_sampling
from medis.utils import dprint
from medis.plot_tools import quick2D

################################################################################
# Deformable Mirror
################################################################################
def deformable_mirror(wf, WFS_map, iter, previous_output=None, apodize=False, plane_name='', debug=False):
    """
    combine different DM actuator commands into single map to send to prop_dm

    prop_dm needs an input map of n_actuators x n_actuators in units of actuator command height. quick_ao will handle
    the conversion to actuator command height, and the CDI probe must be scaled in cdi.probe_amp in params in
    units of m. Each subroutine is also responsible for creating a map of n_actuators x n_actuators spacing. prop_dm
    handles the resampling of this map onto the wavefront, including the influence function. Its some wizardry that
    happens in c, and presumably it is taken care of so you don't have to worry about it.

    In the call to proper.prop_dm, we apply the flag tp.fit_dm, which switches between two 'modes' of proper's DM
    surface fitting. If FALSE, the DM is driven to the heights specified by dm_map, and the influence function will
    act on these heights to define the final surface shape applied to the DM, which may differ substantially from
    the initial heights specified by dm_map. If TRUE, proper will iterate applying the influence function to the
    input heights, and adjust the heights until the difference between the influenced-map and input map meets some
    proper-defined convergence criterea. Setting tp.fit_dm=TRUE will obviously slow down the code, but will (likely)
    more accurately represent a well-calibrated DM response function.

    much of this code copied over from example from Proper manual on pg 94

    :param wf: single wavefront
    :param WFS_map: wavefront sensor map, should be in units of phase delay
    :param previous_output:
    :param iter: the current index of iteration (which timestep this is)
    :param plane_name: name of plane (should be 'woofer' or 'tweeter' for best functionality)
    :return: nothing is returned, but the probe map has been applied to the DM via proper.prop_dm. DM plane post DM
        application can be saved via the sp.save_list functionality
    """
    assert np.logical_xor(WFS_map is None, previous_output is None)

    # AO Actuator Count from DM Type
    if plane_name == 'tweeter' and hasattr(tp,'act_tweeter'):
        nact = tp.act_tweeter
    elif plane_name == 'woofer' and hasattr(tp,'act_woofer'):
        nact = tp.act_woofer
    else:
        nact = tp.ao_act

    # DM Coordinates
    nact_across_pupil = nact - 2  # number of full DM actuators across pupil (oversizing DM extent)
    dm_xc = (nact / 2)   # The location of the optical axis (center of the wavefront) on the DM in
    dm_yc = (nact / 2)   # actuator units. First actuator is centered on (0.0, 0.0). The 0.5 is a
    #  parameter introduced/tuned by Rupert to remove weird errors (address this).
    # KD verified this needs to be here or else suffer weird errors 9/19
    # TODO address/remove the 0.5 in DM x,y coordinates

    ############################
    # Creating DM Surface Map
    ############################
    d_beam = 2 * proper.prop_get_beamradius(wf)  # beam diameter
    act_spacing = d_beam / nact_across_pupil  # actuator spacing [m]

    #######
    # AO
    #######
    if previous_output is not None and WFS_map is None:
        dm_map = update_dm(previous_output)
    else:
        dm_map = quick_ao(wf, nact, WFS_map[wf.iw])

    #########
    # Waffle
    #########
    if tp.satelite_speck['apply'] and plane_name is not 'woofer':
        waffle = make_speckle_kxy(tp.satelite_speck['xloc'], tp.satelite_speck['yloc'],
                                  tp.satelite_speck['amp'], tp.satelite_speck['phase'])
        waffle += make_speckle_kxy(tp.satelite_speck['xloc'], -tp.satelite_speck['yloc'],
                                   tp.satelite_speck['amp'], tp.satelite_speck['phase'])
        dm_map += waffle

    #######
    # CDI
    ######
    if cdi.use_cdi and plane_name == cdi.which_DM:
        theta = cdi.phase_series[iter]
        if not np.isnan(theta):
            # dprint(f"Applying CDI probe, lambda = {wfo.wsamples[iw]*1e9:.2f} nm")
            cdi.save_tseries(iter, datetime.datetime.now())
            probe = config_probe(theta, nact, iw=wf.iw, ib=wf.ib, tstep=iter)
            dm_map = dm_map + probe  # Add Probe to DM map

    #########################
    # Applying Piston Error
    #########################
    if tp.piston_error:
        mean_dm_map = np.mean(np.abs(dm_map))
        var = 1e-4  # 1e-11
        dm_map = dm_map + np.random.normal(0, var, (dm_map.shape[0], dm_map.shape[1]))

    #########################
    # proper.prop_dm
    #########################
    dmap = proper.prop_dm(wf, dm_map, dm_xc, dm_yc, act_spacing, FIT=tp.fit_dm)  #

    if debug and wf.iw == 0 and wf.ib == 0 and iter==0:
        dprint(plane_name)
        check_sampling(wf, iter, plane_name+' DM pupil plane', getframeinfo(stack()[0][0]), units='mm')

        quick2D(WFS_map[wf.iw], title=f"WFS map after masking",
                zlabel='unwrapped phase (rad)',
                vlim=[-3 * np.pi, 3 * np.pi])

        fig, ax = plt.subplots(1,1)
        cax = ax.imshow(dm_map*1e9, interpolation='none', origin='lower')
        plt.title(f'{plane_name} dm_map (actuator coordinates)')
        cb = plt.colorbar(cax)
        cb.set_label('nm')

        plt.show()

        post_ao = unwrap_phase(proper.prop_get_phase(wf)) * wf.lamda / (2 * np.pi)
        # quick2D(pre_ao_dist*1e9, title='unwrapped wavefront before DM', zlabel='nm', show=False)  # , vlim=(-0.5e-7,0.5e-7))
        # quick2D(np.abs(pre_ao_amp)**2, title='Pre-AO Intensity', show=False)#, vlim=(-0.5e-7,0.5e-7))
        # quick2D(dmap, title='the phase map prop_dm is applying', zlabel='distance (m)', show=False)#, vlim=(-0.5e-7,0.5e-7))
        # plt.figure()
        # plt.plot(pre_ao_dist[len(pre_ao_dist)//2], label=f'pre_ao 1D cut, row {len(pre_ao_dist)//2}')
        # plt.plot(2*dmap[len(dmap)//2], label=f'dmap 1D cut (x2), row {len(dmap)//2}')
        # plt.plot((pre_ao_dist + (2*dmap))[len(dmap)//2], label='difference')
        # plt.legend()
        # plt.xlim(sp.grid_size//2*np.array([1-sp.beam_ratio*1.1, 1+sp.beam_ratio*1.1]))
        # quick2D(pre_ao + (2*dmap), title='diff', zlabel='m', show=False, vlim=(-0.5e-7,0.5e-7))
        # quick2D(post_ao, title='unwrapped wavefront after DM', zlabel='m', show=True, vlim=(-0.5e-7,0.5e-7))
        # quick2D(np.abs(proper.prop_get_amplitude(wf))**2, title='wavefront after DM intensity', show=False)
        # quick2D(proper.prop_get_phase(wf), title='wavefront after DM in phase units', zlabel='Phase',
        #          show=True)  # colormap='sunlight',

    if apodize:
        hardmask_pupil(wf)

    return dmap

################################################################################
# Ideal AO
################################################################################
def quick_ao(wf, nact, WFS_map):
    """
    calculate the offset map to send to the DM from the WFS map

    The main idea is to apply the DM only to the region of the wavefront that contains the beam. The phase map from
    the wfs saved the whole wavefront, so that must be cropped. During the wavefront initialization in
    wavefront.initialize_proper, the beam ratio set in sp.beam_ratio is scaled per wavelength (to achieve constant
    sampling sto create white light images), so the cropped value must also be scaled by wavelength. Note, beam ratio
    is scaled differently depending on if sp.focused_sys is True or not. See params-->sp.focused_sys and Proper
    manual pg 36 for more info.

    Then, we interpolate the cropped beam onto a grid of (n_actuators,n_actuators), such that the DM can apply a
    actuator height to each represented actuator, not a over or sub-sampled form. If the number of actuators is low
    compared to the number of samples on the beam, you should anti-alias the WFS map via a lowpass filter before
    interpolating. There is a discrepancy between the sampling of the wavefront at this location (the size you cropped)
    vs the size of the DM. proper.prop_dm handles this, so just plug in the n_actuator sized DM map with specified
    parameters, and assume that prop_dm handles the resampling correctly via the spacing or n_act_across_pupil flag.
    FYI the resampling is done via a c library you installed/compiled when installing proper.

    The WFS map is a map of real values in units of phase delay in radians. However, the AO map that gets passed to
    proper.prop_dm wants input in nm height of each actuator. Therefore, you need to convert the phase delay to
    a DM height. For the ideal AO, you would do this individually for each wavelength. However, for a 'real' AO system
    you do this for the median wavelength. You also need to account for a factor of 2, since the DM is modeled as
    a mirror so it travels the length of the phase delay twice.

    much of this code copied over from example from Proper manual on pg 94

    :param wfo: wavefront object created by optics.Wavefronts() [n_wavelengths, n_objects] of tp.gridsize x tp.gridsize
    :param WFS_map: returned from quick_wfs (as of Aug 2019, its an idealized image)
    :return: ao_map: map of DM actuator command heights in units of m
    """

    nact_across_pupil = nact-2          # number of full DM actuators across pupil (oversizing DM extent)
                                        # Note: oversample by 2 actuators hardcoded here, check if this is appropriate

    ############################
    # Creating AO Surface Map
    ############################
    d_beam = 2 * proper.prop_get_beamradius(wf)  # beam diameter
    act_spacing = d_beam / nact_across_pupil  # actuator spacing [m]

    ###################################
    # Cropping the Beam from WFS map
    ###################################
    # cropping here by beam_ratio rather than d_beam is valid since the beam size was initialized
    #  using the scaled beam_ratios when the wfo was created
    # crop should be -1,+1 on either side of the center because for an even sp.grid_size
    ao_map = WFS_map[
             sp.grid_size//2 - np.int_(wf.beam_ratio*sp.grid_size//2)-1:
             sp.grid_size//2 + np.int_(wf.beam_ratio*sp.grid_size//2)+2,
             sp.grid_size//2 - np.int_(wf.beam_ratio*sp.grid_size//2)-1:
             sp.grid_size//2 + np.int_(wf.beam_ratio*sp.grid_size//2)+2]
    # dprint(f"WFS map coordinates are {sp.grid_size//2 - np.int_(wf.beam_ratio*sp.grid_size//2)-1},"
    #        f"{sp.grid_size//2 + np.int_(wf.beam_ratio*sp.grid_size//2)+1}")

    ########################################################
    # Interpolating the WFS map onto the actuator spacing
    # (tp.nact,tp.nact)
    ########################################################
    # Lowpass Filter- prevents aliasing; uses Gaussian filter
    nyquist_dm = nact/2 * act_spacing  # [m]
    sigma = [nyquist_dm/2.355, nyquist_dm/2.355]  # assume we want sigma to be twice the HWHM
    ao_map = ndimage.gaussian_filter(ao_map, sigma=sigma, mode='nearest')

    f = interpolate.interp2d(range(ao_map.shape[0]), range(ao_map.shape[0]), ao_map, kind='cubic')
    ao_map = f(np.linspace(0,ao_map.shape[0],nact), np.linspace(0,ao_map.shape[0], nact))
    # map_spacing = proper.prop_get_sampling(wf)
    # ao_map = proper.prop_magnify(ao_map, map_spacing / act_spacing, nact, QUICK=True)

    ################################################
    # Converting phase delay to DM actuator height
    ################################################
    # Apply the inverse of the WFS image to the DM, so use -dm_map (dm_map is in phase units, divide by k=2pi/lambda)
    surf_height = proper.prop_get_wavelength(wf) / (4 * np.pi)  # [m/rad]
    ao_map = -ao_map * surf_height  # Converts DM map to units of [m] of actuator heights

    return ao_map

def retro_wfs(star_fields, wfo, plane_name='wfs'):
    """
    Retrospective wfs (measure an old field)

    :param star_fields:
    :param wfo:
    :param plane_name:
    :return:
    """
    WFS_map = np.zeros((len(star_fields), sp.grid_size, sp.grid_size))
    from skimage.restoration import unwrap_phase
    for iw in range(len(star_fields)):
        quick2D(np.angle(star_fields), title='before mask', colormap='sunlight')
        phasemap = np.angle(star_fields[iw])
        masked_phase = np.ma.masked_equal(phasemap, 0)
        quick2D(masked_phase, title='before unwrap', colormap='sunlight')
        WFS_map[iw] = unwrap_phase(masked_phase, wrap_around=[False, False])
        WFS_map[iw][phasemap == 0] = 0
        quick2D(WFS_map[iw], title='after')
    if 'retro_closed_wfs' in sp.save_list:
        wfo.save_plane(location='WFS_map')

    return WFS_map

def open_loop_wfs(wfo, plane_name='wfs'):
    """
    saves the unwrapped phase [arctan2(imag/real)] of the wfo.wf_collection at each wavelength

    It is an idealized image (exact copy) of the wavefront phase per wavelength. Only the map for the first object
    (the star) is saved. We have initialized

    Here we hardmask on the WFS map to be a circle around the beam in the pupil plane. This hard masking prevents the
     DM from acting on non-beam signal, since the DM modelled by proper is a nxn square array, but the beam is nominally
     circular for circular apertures.

    #TODO the way this is saved for naming the WFS_map is going to break if you want to do closed loop WFS on a
    #TODO woofer-tweeter system

    :param wfo: wavefront object
    :param plane_name: name of the plane to enable or disable saving the WFS map
    :return: array containing only the unwrapped phase delay of the wavefront; shape=[n_wavelengths], units=radians
    """
    star_wf = wfo.wf_collection[:, 0]
    WFS_map = np.zeros((len(star_wf), sp.grid_size, sp.grid_size))

    for iw in range(len(star_wf)):  # for each wavelength
        hardmask_pupil(star_wf[iw])
        phasemap = proper.prop_get_phase(star_wf[iw])
        WFS_map[iw] = unwrap_phase(phasemap, wrap_around=[False, False])
        WFS_map[iw][phasemap==0] = 0 #TODO is this still necessary?

        # if sp.verbose:
        #     quick2D(WFS_map[iw], title=f"WFS map after masking, lambda={wfo.wsamples[iw]*1e9:.2f}",
        #             zlabel='unwrapped phase (rad)',
        #             vlim=[-3*np.pi, 3*np.pi])
        #

    if 'WFS' in sp.save_list or sp.closed_loop:
        wfo.save_plane(location='WFS')

    return WFS_map

def hardmask_pupil(wf):
    """
    hard-edged circular mask of the pupil plane.

    Masks out the WFS map outside of the beam since the DM modeled by proper can only be a square nxn grid of actuators,
    and thus the influence function surrounding each DM actuator could  be affecting on-beam pixels, even if the
    actuator is acting on off-beam signal. In other words, even if a cornerDM actuator doesn't actuate on the beam,
    if there was non-zero signal in the WFS map, it will try to act on it, and it could 'influence' nearby DM actuators
    that are acting on the beam.

    This hard-edged mask is different from prop_circular_aperture in that it does not anti-alias the edges of the mask
    based on the 'fill factor' of the edge pixels. Instead, it has a boolean mask to zero everything > a fixed radius,
    in this case determined by the grid size and beam ratio of each wavefront passed into it.

    :param wf: a single wavefront
    :return: nothing is returned but the wf passed into it has been masked

    """
    phase_map = proper.prop_get_phase(wf)
    amp_map = proper.prop_get_amplitude(wf)

    # Sizing the Mask
    h, w = wf.wfarr.shape[:2]
    center = (int(w / 2), int(h / 2))
    radius = np.floor(sp.grid_size * wf.beam_ratio / 2)  # Should scale with wavelength if sp.focused_system=False,
                                                        # np.ceil used to oversize map so don't clip the beam
    # Making the Circular Boolean Mask
    Y, X = np.mgrid[:h, :w]
    dist_from_center = np.sqrt((X - center[0]) ** 2 + (Y - center[1]) ** 2)
    inds = dist_from_center <= radius

    # Applying the Mask to the Complex Array
    mask = np.zeros_like(phase_map)
    mask[inds] = 1
    masked = phase_map * mask
    wf.wfarr = proper.prop_shift_center(amp_map * np.cos(masked) + 1j * amp_map * np.sin(masked))

    # if sp.verbose:
    #     dprint(f"Radius of hard-edge pupil mask is {radius} pixels")
    #     quick2D(masked, title=f"Masked phase map in hardmask_pupil, lambda={wf.lamda*1e9} nm", zlabel='phase (rad)')
    #     plt.show()

def make_speckle_kxy(kx, ky, amp, dm_phase):
    """given an kx and ky wavevector,
    generates a NxN flatmap that has
    a speckle at that position"""
    N = tp.ao_act
    dmx, dmy   = np.meshgrid(
                    np.linspace(-0.5, 0.5, N),
                    np.linspace(-0.5, 0.5, N))

    xm=dmx*kx*2.0*np.pi
    ym=dmy*ky*2.0*np.pi
    # print 'DM phase', dm_phase
    ret = amp*np.cos(xm + ym +  dm_phase)
    return ret

# def ao(wf, WFS_map, theta):
#     if sp.closed_loop:
#         deformable_mirror(wf, WFS_map, theta)
#     else:
#         WFS_map = open_loop_wfs(wf)  # overwrite WFS_map
#         # dprint(f"WFS_ma.shape = {WFS_map.shape}")
#         deformable_mirror(wf, WFS_map, theta)
#
################################################################################
# Full AO
################################################################################
# not implemented. Full AO implies a time delay, and maybe non-ideal WFS


#! /usr/bin/env python

# Eureka! Stage 3 reduction pipeline

# Proposed Steps
# --------------
# 1.  Read in all data frames and header info from Stage 2 data products DONE
# 2.  Record JD and other relevant header information DONE
# 3.  Apply light-time correction (if necessary) DONE
# 4.  Calculate trace and 1D+2D wavelength solutions (if necessary)
# 5.  Make flats, apply flat field correction (Stage 2)
# 6.  Manually mask regions DONE
# 7.  Compute difference frames OR slopes (Stage 1)
# 8.  Perform outlier rejection of BG region DONE
# 9.  Background subtraction DONE
# 10. Compute 2D drift, apply rough (integer-pixel) correction
# 11. Full-frame outlier rejection for time-series stack of NDRs
# 12. Apply sub-pixel 2D drift correction
# 13. Extract spectrum through summation DONE
# 14. Compute median frame DONE
# 15. Optimal spectral extraction DONE
# 16. Save Stage 3 data products
# 17. Produce plots DONE

import os, time, glob
import numpy as np
from astropy.io import fits
from tqdm import tqdm
from . import optspex
from . import plots_s3, source_pos
from . import background as bg
from . import bright2flux as b2f
from ..lib import sort_nicely as sn
from ..lib import logedit
from ..lib import readECF as rd
from ..lib import manageevent as me
from ..lib import astropytable
from ..lib import util

class MetaClass:
    '''A class to hold Eureka! metadata.
    '''

    def __init__(self):
        return

class DataClass:
    '''A class to hold Eureka! image data.
    '''

    def __init__(self):
        return

def reduceJWST(eventlabel, s2_meta=None):
    '''Reduces data images and calculates optimal spectra.

    Parameters
    ----------
    eventlabel: str
        The unique identifier for these data.
    s2_meta:    MetaClass
        The metadata object from Eureka!'s S2 step (if running S2 and S3 sequentially).

    Returns
    -------
    meta:   MetaClass
        The metadata object with attributes added by S3.

    Notes
    -------
    History:

    - May 2021 Kevin Stevenson
        Initial version
    - October 2021 Taylor Bell
        Updated to allow for inputs from S2
    '''

    # Initialize data object
    data = DataClass()

    # Initialize a new metadata object
    meta = MetaClass()
    meta.eventlabel = eventlabel

    # Load Eureka! control file and store values in Event object
    ecffile = 'S3_' + eventlabel + '.ecf'
    ecf = rd.read_ecf(ecffile)
    rd.store_ecf(meta, ecf)
    meta.eventlabel=eventlabel

    # S3 is not being called right after S2 - try to load a metadata in case S2 was previously run
    if s2_meta == None:
        # Search for the S2 output metadata in the inputdir provided in
        rootdir = os.path.join(meta.topdir, *meta.inputdir.split(os.sep))
        if rootdir[-1]!='/':
            rootdir += '/'
        fnames = glob.glob(rootdir+'**/S2_'+meta.eventlabel+'_Meta_Save.dat', recursive=True)
        fnames = sn.sort_nicely(fnames)

        if len(fnames)==0:
            # There may be no metafiles in the inputdir - raise an error and give a helpful message
            print('WARNING: Unable to find an output metadata file from Eureka!\'s S2 step '
                 +'in the inputdir: \n"{}"!\n'.format(meta.inputdir)
                 +'Assuming this S2 data was produced by the JWST pipeline instead.')
        else:
            if len(fnames)>1:
                # There may be multiple runs - use the most recent but warn the user
                print('WARNING: There are multiple metadata save files in your inputdir: \n"{}"\n'.format(meta.inputdir)
                     +'Using the metadata file: \n"{}"'.format(fnames[-1]))

            fname = fnames[-1] # Pick the last file name
            fname = fname[:-4] # Strip off the .dat ending

            s2_meta = me.loadevent(fname)

    # Locate the exact output folder from the previous S2 run (since there is a procedurally generated subdirectory for each run)
    if s2_meta != None:
        # Need to remove the topdir from the outputdir
        if os.path.isdir(s2_meta.outputdir):
            s2_outputdir = s2_meta.outputdir[len(s2_meta.topdir):]
            if s2_outputdir[0]=='/':
                s2_outputdir = s2_outputdir[1:]

            meta = s2_meta

            # Load Eureka! control file and store values in the S2 metadata object
            ecffile = 'S3_' + eventlabel + '.ecf'
            ecf = rd.read_ecf(ecffile)
            rd.store_ecf(meta, ecf)

            # Overwrite the inputdir with the exact output directory from S2
            meta.inputdir = s2_outputdir
        else:
            raise AssertionError("Unable to find output data files from Eureka!'s S2 step! "
                                 + "Looked in the folder: \n{}".format(s2_meta.outputdir))

    meta.inputdir_raw = meta.inputdir
    meta.outputdir_raw = meta.outputdir

    # check for range of spectral apertures
    if isinstance(meta.spec_hw, list):
        meta.spec_hw_range = range(meta.spec_hw[0], meta.spec_hw[1]+meta.spec_hw[2], meta.spec_hw[2])
    else:
        meta.spec_hw_range = [meta.spec_hw]

    #check for range of background apertures
    if isinstance(meta.bg_hw, list):
        meta.bg_hw_range = range(meta.bg_hw[0], meta.bg_hw[1]+meta.spec_hw[2], meta.bg_hw[2])
    else:
        meta.bg_hw_range = [meta.bg_hw]

    # create directories to store data
    meta.runs = [] # Used to make sure we're always looking at the right run for each aperture/annulus pair
    for spec_hw_val in meta.spec_hw_range:

        for bg_hw_val in meta.bg_hw_range:

            meta.eventlabel = eventlabel

            meta.runs.append(util.makedirectory(meta, 'S3', ap=spec_hw_val, bg=bg_hw_val))

    # begin process
    run_i = 0
    for spec_hw_val in meta.spec_hw_range:

        for bg_hw_val in meta.bg_hw_range:

            t0 = time.time()

            meta.spec_hw = spec_hw_val

            meta.bg_hw = bg_hw_val

            meta.outputdir = util.pathdirectory(meta, 'S3', meta.runs[run_i], ap=spec_hw_val, bg=bg_hw_val)
            run_i += 1

            event_ap_bg = meta.eventlabel + "_ap" + str(spec_hw_val) + '_bg' + str(bg_hw_val)

            # Open new log file
            meta.logname = meta.outputdir + 'S3_' + event_ap_bg + ".log"
            log = logedit.Logedit(meta.logname)
            log.writelog("\nStarting Stage 3 Reduction\n")
            log.writelog(f"Input directory: {meta.inputdir}")
            log.writelog(f"Output directory: {meta.outputdir}")
            log.writelog("Using ap=" + str(spec_hw_val) + ", bg=" + str(bg_hw_val))

            # Copy ecf (and update inputdir in case S3 is being called sequentially with S2)
            log.writelog('Copying S3 control file')
            new_ecfname = meta.outputdir + ecffile.split('/')[-1]
            with open(new_ecfname, 'w') as new_file:
                with open(ecffile, 'r') as file:
                    for line in file.readlines():
                        if len(line.strip())==0 or line.strip()[0]=='#':
                            new_file.write(line)
                        else:
                            line_segs = line.strip().split()
                            if line_segs[0]=='inputdir':
                                new_file.write(line_segs[0]+'\t\t/'+meta.inputdir+'\t'+' '.join(line_segs[2:])+'\n')
                            else:
                                new_file.write(line)

            # Create list of file segments
            meta = util.readfiles(meta)
            num_data_files = len(meta.segment_list)
            if num_data_files==0:
                rootdir = os.path.join(meta.topdir, *meta.inputdir.split(os.sep))
                if rootdir[-1]!='/':
                    rootdir += '/'
                raise AssertionError(f'Unable to find any "{meta.suffix}.fits" files in the inputdir: \n"{rootdir}"!')
            else:
                log.writelog(f'\nFound {num_data_files} data file(s) ending in {meta.suffix}.fits')

            with fits.open(meta.segment_list[-1]) as hdulist:
                # Figure out which instrument we are using
                meta.inst = hdulist[0].header['INSTRUME'].lower()
            # Load instrument module
            if meta.inst == 'miri':
                from . import miri as inst
            elif meta.inst == 'nircam':
                from . import nircam as inst
            elif meta.inst == 'nirspec':
                from . import nirspec as inst
            elif meta.inst == 'niriss':
                raise ValueError('NIRISS observations are currently unsupported!')
            else:
                raise ValueError('Unknown instrument {}'.format(meta.inst))

            stdspec = np.array([])
            # Loop over each segment
            # Only reduce the last segment/file if testing_S3 is set to True in ecf
            if meta.testing_S3:
                istart = num_data_files - 1
            else:
                istart = 0
            for m in range(istart, num_data_files):
                # Keep track if this is the first file - otherwise MIRI will keep swapping x and y windows
                if m==istart:
                    meta.firstFile = True
                else:
                    meta.firstFile = False
                # Report progress
                log.writelog(f'Reading file {m + 1} of {num_data_files}')
                # Read in data frame and header
                data, meta = inst.read(meta.segment_list[m], data, meta)
                # Get number of integrations and frame dimensions
                meta.n_int, meta.ny, meta.nx = data.data.shape
                if meta.testing_S3:
                    # Only process the last 5 integrations when testing
                    meta.int_start = np.max((0,meta.n_int-5))
                else:
                    meta.int_start = 0
                # Locate source postion
                meta.src_ypos = source_pos.source_pos(data, meta, m, header=('SRCYPOS' in data.shdr))
                log.writelog(f'  Source position on detector is row {meta.src_ypos}.')
                # Trim data to subarray region of interest
                data, meta = util.trim(data, meta)
                # Create bad pixel mask (1 = good, 0 = bad)
                # FINDME: Will want to use DQ array in the future to flag certain pixels
                data.submask = np.ones(data.subdata.shape)

                # Convert flux units to electrons (eg. MJy/sr -> DN -> Electrons)
                data, meta = b2f.convert_to_e(data, meta, log)

                # Check if arrays have NaNs
                data.submask = util.check_nans(data.subdata, data.submask, log, name='SUBDATA')
                data.submask = util.check_nans(data.suberr, data.submask, log, name='SUBERR')
                data.submask = util.check_nans(data.subv0, data.submask, log, name='SUBV0')

                # Manually mask regions [colstart, colend, rowstart, rowend]
                if hasattr(meta, 'manmask'):
                    log.writelog("  Masking manually identified bad pixels")
                    for i in range(len(meta.manmask)):
                        ind, colstart, colend, rowstart, rowend = meta.manmask[i]
                        data.submask[rowstart:rowend, colstart:colend] = 0

                # Perform outlier rejection of sky background along time axis
                log.writelog('  Performing background outlier rejection')
                meta.bg_y2 = int(meta.src_ypos + bg_hw_val)
                meta.bg_y1 = int(meta.src_ypos - bg_hw_val)
                data = inst.flag_bg(data, meta)

                data = bg.BGsubtraction(data, meta, log, meta.isplots_S3)

                # Calulate drift2D
                # print("Calculating 2D drift...")

                # print("Performing rough, pixel-scale drift correction...")

                # Outlier rejection of full frame along time axis
                # print("Performing full-frame outlier rejection...")

                if meta.isplots_S3 >= 3:
                    log.writelog('  Creating figures for background subtraction')
                    for n in tqdm(range(meta.int_start,meta.n_int)):
                        # make image+background plots
                        plots_s3.image_and_background(data, meta, n)

                # print("Performing sub-pixel drift correction...")

                # Select only aperture region
                ap_y1 = int(meta.src_ypos - spec_hw_val)
                ap_y2 = int(meta.src_ypos + spec_hw_val)
                data.apdata  = data.subdata[:, ap_y1:ap_y2]
                data.aperr   = data.suberr[:, ap_y1:ap_y2]
                data.apmask  = data.submask[:, ap_y1:ap_y2]
                data.apbg    = data.subbg[:, ap_y1:ap_y2]
                data.apv0    = data.subv0[:, ap_y1:ap_y2]
                # Extract standard spectrum and its variance
                data.stdspec = np.sum(data.apdata, axis=1)
                data.stdvar  = np.sum(data.aperr ** 2, axis=1)  # FINDME: stdvar >> stdspec, which is a problem
                # Compute fraction of masked pixels within regular spectral extraction window
                # numpixels   = 2.*meta.spec_width*subnx
                # fracMaskReg = (numpixels - np.sum(apmask,axis=(2,3)))/numpixels

                # Compute median frame
                data.medsubdata = np.median(data.subdata, axis=0)
                data.medapdata  = np.median(data.apdata, axis=0)

                # Extract optimal spectrum with uncertainties
                log.writelog("  Performing optimal spectral extraction")
                data.optspec = np.zeros(data.stdspec.shape)
                data.opterr  = np.zeros(data.stdspec.shape)
                gain = 1  # Already converted DN to electrons, so gain = 1 for optspex
                for n in tqdm(range(meta.int_start,meta.n_int)):
                    data.optspec[n], data.opterr[n], mask = optspex.optimize(data.apdata[n], data.apmask[n], data.apbg[n],
                                                                             data.stdspec[n], gain, data.apv0[n],
                                                                             p5thresh=meta.p5thresh, p7thresh=meta.p7thresh,
                                                                             fittype=meta.fittype, window_len=meta.window_len,
                                                                             deg=meta.prof_deg, n=data.intstart + n,
                                                                             isplots=meta.isplots_S3, eventdir=meta.outputdir,
                                                                             meddata=data.medapdata, hide_plots=meta.hide_plots)

                # Plot results
                if meta.isplots_S3 >= 3:
                    log.writelog('  Creating figures for optimal spectral extraction')
                    for n in tqdm(range(meta.int_start,meta.n_int)):
                        # make optimal spectrum plot
                        plots_s3.optimal_spectrum(data, meta, n)

                # Append results
                if len(stdspec) == 0:
                    wave_2d = data.subwave
                    wave_1d = data.subwave[meta.src_ypos]
                    stdspec = data.stdspec
                    stdvar  = data.stdvar
                    optspec = data.optspec
                    opterr  = data.opterr
                    bjdtdb  = data.bjdtdb
                else:
                    stdspec = np.append(stdspec, data.stdspec, axis=0)
                    stdvar  = np.append(stdvar, data.stdvar, axis=0)
                    optspec = np.append(optspec, data.optspec, axis=0)
                    opterr  = np.append(opterr, data.opterr, axis=0)
                    bjdtdb  = np.append(bjdtdb, data.bjdtdb, axis=0)

            # Calculate total time
            total = (time.time() - t0) / 60.
            log.writelog('\nTotal time (min): ' + str(np.round(total, 2)))

            if meta.testing_S3 == False:
                log.writelog('Saving results as astropy table')
                meta.tab_filename = meta.outputdir + 'S3_' + event_ap_bg + "_Table_Save.txt"
                astropytable.savetable_S3(meta.tab_filename, bjdtdb, wave_1d, stdspec, stdvar, optspec, opterr)

            if meta.isplots_S3 >= 1:
                log.writelog('Generating figure')
                # 2D light curve without drift correction
                plots_s3.lc_nodriftcorr(meta, wave_1d, optspec)

            # Save results
            if meta.testing_S3 == False:
                log.writelog('Saving Metadata')
                me.saveevent(meta, meta.outputdir + 'S3_' + event_ap_bg + "_Meta_Save", save=[])

            log.closelog()

    return meta

import numpy,import astropy,from astropy


"""The Mine class."""

__docformat__ = "reStructuredText"

from functools import reduce
from typing import List, Dict, Optional, Union, Callable, Tuple
from abc import ABC, abstractmethod
from collections import defaultdict

import numpy as np
import matplotlib.pyplot as plt
import networkx as nx
from prettytable import PrettyTable
from qiskit.opflow import PauliOp, I, Z, Plus, MatrixOp
from qiskit.algorithms import (
    MinimumEigensolver,
    AlgorithmResult,
    NumPyMinimumEigensolver,
)

from ..algorithms import Pseudoflow, PEPSITE
from ..benchmark import Benchmark
from ..utils import null_operator, z_projector, single_qubit_pauli, int_to_bitstr

class BaseMine(ABC):
    """
    Abstract Mine class
    """

    def __init__(self) -> None:
        pass

    @abstractmethod
    def solve(self) -> "MiningProblemResult":
        raise NotImplementedError()

class Mine(BaseMine):
    """This class stores the mine configurations."""

    def __init__(self, mine_config: Union[str, np.ndarray]) -> None:
        """Initialize the Mine class.

        Parameters
        ----------
        mine_config : str
            Path to the mine configuration file.

        Raises
        ------
        IOError
            Invalid path to the mine configuration file.

        """
        if isinstance(mine_config, str):
            try:
                self.dat = np.loadtxt(mine_config, dtype=float)
            except:
                raise IOError("Invalid Mine Configuration File")
        elif isinstance(mine_config, np.ndarray):
            if len(mine_config.shape) != 2:
                raise ValueError("`mine_config` must be two-demensional")
            self.dat = np.array(mine_config, dtype=float)
        else:
            raise ValueError("Unrecognized `mine_config` type")

        self.rows, self.cols = self.dat.shape
        self.graph = defaultdict(list)  # p:c
        self.graph_r = defaultdict(list)  # c:p
        self.idx2cord = []
        self.cord2idx = {}
        self._init_mapping()
        self.nqubits: int = len(self.idx2cord)
        self.valid_configs = None
        self._Hs = None
        self._Hp = None

    @property
    def Hs(self) -> float:
        if self._Hs is None:
            self._Hs = self.gen_Hs()
        return self._Hs

    @property
    def Hp(self) -> float:
        if self._Hp is None:
            self._Hp = self.gen_Hp()
        return self._Hp

    @staticmethod
    def gen_random_mine(
        size: Tuple[int, int], distribution: Optional[str] = "normal", seed: Optional[int] = None, **kwargs
    ) -> "Mine":
        distribution = distribution.lower()
        rng = np.random.default_rng(seed)
        if distribution == "gaussian" or distribution == "normal":
            return Mine(rng.normal(size=size, **kwargs))
        if distribution == "uniform":
            return Mine(rng.uniform(size=size, **kwargs))
        raise ValueError("Distribution not supported.")

    def _init_mapping(self) -> None:
        """Assign a unique id to each valid node and store the graph structure with ``self.graph``."""
        for r in range(self.rows):
            for c in range(self.cols):
                if c < r or c > self.cols - r - 1:
                    self.dat[r, c] = float("inf")
                if self.dat[r, c] < float("inf"):
                    self.idx2cord.append((r, c))
                    self.cord2idx[(r, c)] = len(self.idx2cord) - 1
                    idx = self.cord2idx[(r, c)]
                    for pr, pc in [(r - 1, c - 1), (r - 1, c), (r - 1, c + 1)]:
                        if (
                            0 <= pr < self.rows
                            and 0 <= pc < self.cols
                            and self.dat[pr, pc] < float("inf")
                        ):
                            self.graph[idx].append(self.cord2idx[(pr, pc)])
                            self.graph_r[self.cord2idx[(pr, pc)]].append(idx)

    def plot_mine(self) -> None:
        """Plot the mine configuration."""
        x = PrettyTable([" "] + [str(ic) for ic in range(self.cols)])
        for ir in range(self.rows):
            x.add_row([ir] + ["%.3f" % self.dat[ir, ic] for ic in range(self.cols)])
        print(str(x))

    def plot_mine_graph(
        self, color: str = "r", pos_func: Callable = nx.spring_layout
    ) -> None:
        """Plot a graph representing the mine configuration.

        Parameters
        ----------
        color : str
            Color of the nodes.
        pos_func : Callable
            A Callable that returns positions of nodes.

        """
        G = nx.Graph()
        G.add_nodes_from(np.arange(self.nqubits))
        elist = [[i, j] for i, jl in self.graph.items() for j in jl]
        # tuple is (i,j,weight) where (i,j) is the edge
        G.add_edges_from(elist)

        colors = [color for node in G.nodes()]
        pos = pos_func(G)

        default_axes = plt.axes(frameon=True)
        nx.draw_networkx(
            G, node_color=colors, node_size=600, alpha=0.8, ax=default_axes, pos=pos
        )

    def gen_Hs(self) -> PauliOp:
        """Generate the smoothess Hamiltonian
        :math:`H_{s}=\sum_{i}\sum_{j:Parent(i)} (1-Z_{i})/2*(1+Z_{j})/2`

        Returns
        ----------
        qiskit.opflow.PauliOp
            Smoothness Hamiltonian.

        """
        Hs = 0 * I ^ self.nqubits
        for i in range(self.nqubits):
            for j in self.graph[i]:
                Hs += (
                    (I ^ self.nqubits) - ((I ^ self.nqubits - i - 1) ^ Z ^ (I ^ i))
                ) @ ((I ^ self.nqubits) + ((I ^ self.nqubits - j - 1) ^ Z ^ (I ^ j)))
        return 0.25 * Hs

    def gen_Hp(self) -> PauliOp:
        """Generate the profit Hamiltonian
        :math:`H_{p}=\sum_{i}w(i)(1-Z_{i})/2`

        Returns
        ----------
        qiskit.opflow.PauliOp
            Profit Hamiltonian.

        """
        Hp = 0 * I ^ self.nqubits
        for i in range(self.nqubits):
            # Qiskit opflow doesn't support multiplying types other than python built-in ones (e.g. np.float64)
            Hp += float(self.dat[self.idx2cord[i]]) * (
                (I ^ self.nqubits) - ((I ^ self.nqubits - i - 1) ^ Z ^ (I ^ i))
            )
        return 0.5 * Hp

    def gen_Hamiltonian(self, penalty: Union[float, bool, None]) -> PauliOp:
        """Generate the Hamiltonian with penalty weight :math:`\gamma`.

        :math:`H=-H_{p}+\gamma H_{s}`

        Parameter   s
        ----------
        penalty : float
            Penalty for the smoothness term in the Hamiltonian.

        Returns
        ----------
        qiskit.opflow.PauliOp
            Hamiltonian with penalty weight :math:`\gamma`.

        """
        if penalty is True:
            penalty = self.heuristic_penalty()
        if penalty:
            return (-self.Hp + penalty * self.Hs).reduce()
        return self.gen_projected_Hamiltonian()

    def gen_projected_Hamiltonian(self) -> PauliOp:
        """Generate the profit Hamiltonian projected onto valid states.

        Returns
        ----------
        qiskit.opflow.PauliOp
            :math:`H=-PH_{p}P`

        """
        state_fn = (-self.Hs @ (Plus ^ self.nqubits)).reduce().eval().to_dict_fn()
        self.valid_configs = [
            int(k, 2) for k, v in state_fn.primitive.items() if abs(v) < 1e-8
        ]
        p_op = np.zeros((2 ** self.nqubits))
        p_op[np.array(self.valid_configs, dtype=int)] = 1
        p_op = MatrixOp(np.diag(p_op))

        return (p_op @ -self.Hp @ p_op).reduce().to_matrix_op()

    def gen_pseudoflow_graph(self, MAX_FLOW: int) -> Tuple[nx.DiGraph, int, int]:
        G = nx.DiGraph()
        G.add_nodes_from(np.arange(self.nqubits))
        source = -1
        sink = self.nqubits

        for p in self.graph:
            for c in self.graph[p]:
                G.add_edge(p, c, const=MAX_FLOW)

            if self.dat[self.idx2cord[p]] >= 0:
                G.add_edge(source, p, const=self.dat[self.idx2cord[p]])
            else:
                G.add_edge(p, sink, const=-self.dat[self.idx2cord[p]])

        return G, source, sink

    def plot_mine_state(
        self, bitstring: str, bit_ordering: Optional[str] = "R"
    ) -> None:
        """Plot the mining state represented by the bitstring.

        Parameters
        ----------
        bitstring : str
            A 0/1 string represents the state. Length of the string is the same as the number of qubits.
        bit_ordering : str
            Available options ``[\'L\', \'R\']``. ``\'L\'`` means the least significant bit (LSB) is on the left, and ``\'R\'`` means LSB is on the right. LSB represents the qubit with index 0.
        """
        assert (
            len(bitstring) == self.nqubits
        ), "Length of the bitstring should be the same as the number of qubits."
        assert bit_ordering in ["L", "R"], "bit_ordering options: 'L', 'R'."
        if bit_ordering == "R":
            bitstring = "".join(list(bitstring)[::-1])

        x = PrettyTable([" "] + [str(ic) for ic in range(self.cols)])
        for ir in range(self.rows):
            x.add_row(
                [ir]
                + [
                    bitstring[self.cord2idx[(ir, ic)]]
                    if (ir, ic) in self.cord2idx
                    else "x"
                    for ic in range(self.cols)
                ]
            )
        print(str(x))

    def get_profit(self, bitstring: str, bit_ordering: Optional[str] = "R") -> int:
        """Return profit for a vector state.

        Parameters
        ----------
        bitstring : str
            A 0/1 string represents the state. Length of the string is the same as the number of qubits.
        bit_ordering : str
            Available options ``[\'L\', \'R\']``. ``\'L\'`` means the least significant bit (LSB) is on the left, and ``\'R\'`` means LSB is on the right. LSB represents the qubit with index 0.
        """
        assert (
            len(bitstring) == self.nqubits
        ), "Length of the bitstring should be the same as the number of qubits."
        assert bit_ordering in ["L", "R"], "bit_ordering options: 'L', 'R'."
        if bit_ordering == "R":
            bitstring = "".join(list(bitstring)[::-1])

        return sum(
            [
                self.dat[self.idx2cord[i]]
                for i in range(self.nqubits)
                if bitstring[i] == "1"
            ]
        )

    def get_violation(self, bitstring: str, bit_ordering: Optional[str] = "R") -> int:
        """Return violation for a vector state.

        Parameters
        ----------
        bitstring : str
            A 0/1 string represents the state. Length of the string is the same as the number of qubits.
        bit_ordering : str
            Available options ``[\'L\', \'R\']``. ``\'L\'`` means the least significant bit (LSB) is on the left, and ``\'R\'`` means LSB is on the right. LSB represents the qubit with index 0.
        """
        assert (
            len(bitstring) == self.nqubits
        ), "Length of the bitstring should be the same as the number of qubits."
        assert bit_ordering in ["L", "R"], "bit_ordering options: 'L', 'R'."
        if bit_ordering == "R":
            bitstring = "".join(list(bitstring)[::-1])

        dig = list(map(lambda x: -1 if x == "1" else 1, list(bitstring)))
        res = 0
        for i in range(self.nqubits):
            for j in self.graph[i]:
                res += 0.25 * (1.0 - dig[i]) * (1.0 + dig[j])
        return int(res)

    def heuristic_penalty(self, coeff: float = 3.8) -> float:
        return (
            float(
                np.linalg.norm(np.where(self.dat.flat != np.inf), ord=2) / self.nqubits
            )
            * coeff
        )

    def solve(
        self,
        algorithm: Union[MinimumEigensolver, Pseudoflow],
        penalty: Union[float, bool, None] = None,
        benchmark: Union[Benchmark, bool, None] = None,
    ) -> "MiningProblemResult":
        if benchmark is True:
            benchmark = Benchmark()
        elif benchmark is False or benchmark is None:
            benchmark = Benchmark(activate=False)

        with benchmark:
            self._ret = MiningProblemResult()

            if isinstance(algorithm, MinimumEigensolver):
                res = algorithm.compute_minimum_eigenvalue(
                    self.gen_Hamiltonian(penalty), [self.Hp, self.Hs]
                )

                if isinstance(algorithm, NumPyMinimumEigensolver):
                    self._ret.optimal_config = format(
                        np.argwhere(res.eigenstate.primitive.data.real == 1.0).item(),
                        f"0{self.nqubits + 2}b",
                    )[2:]
                    self._ret.optimal_config_prob = 1.0

                else:
                    if isinstance(res.eigenstate, dict):
                        self._ret.optimal_config, self._ret.optimal_config_prob = max(
                            (
                                item
                                for item in res.eigenstate.items()
                                if self.valid_configs is None
                                or int(item[0], 2) in self.valid_configs
                            ),
                            key=lambda item: item[1],
                        )
                    elif isinstance(res.eigenstate, np.ndarray):
                        idx = np.argmax(res.eigenstate * res.eigenstate.conj())
                        self._ret.optimal_config = int_to_bitstr(idx, self.nqubits)
                        self._ret.optimal_config_prob = res.eigenstate[idx]

                    if isinstance(self._ret.optimal_config_prob, complex):
                        self._ret.optimal_config_prob = (
                            self._ret.optimal_config_prob.real ** 2
                            - self._ret.optimal_config_prob.imag ** 2
                        )
                    else:
                        self._ret.optimal_config_prob **= 2

                self._ret.ground_state = res.eigenstate
                if (
                    res.aux_operator_eigenvalues is not None
                ):  # Remove after implemented ASP expectation value
                    (
                        self._ret._expected_profit,
                        self._ret._expected_violation,
                    ) = res.aux_operator_eigenvalues

            elif isinstance(algorithm, Pseudoflow):
                self._ret.optimal_config = algorithm.run(
                    *self.gen_pseudoflow_graph(algorithm.MAX_FLOW)
                )
                self._ret.optimal_config_prob = 1.0

            elif isinstance(algorithm, PEPSITE):
                self._ret.optimal_config = algorithm.run(self)
                self._ret.optimal_config_prob = 1.0

            else:
                raise ValueError(f"{type(algorithm)} is not a valid algorithm.")

            print(
                "The most probable configuration and the corresponding probability:"
                f" {self._ret.optimal_config, self._ret.optimal_config_prob}"
            )
            self.plot_mine_state(self._ret.optimal_config)
            return self._ret

class SubMine(BaseMine):
    def __init__(self, mine: Mine, node_list: List) -> None:
        self.node_list = node_list
        self.nqubits = len(self.node_list)
        self.node_c = {}  # key: a node in this frag, value: child node in another frag
        self.node_p = {}  # key: a node in this frag, value: parent node in another frag
        self.node_bd = []  # submine idx of bd nodes
        self.zb = []  # pauli z op of bd nodes

        self._init_nodes(mine)
        self.Hp = self.gen_Hp(mine)  # profit
        self.Hs = self.gen_Hs(mine)  # smoothness, inner nodes
        self.Hb = null_operator(self.nqubits)  # boundary

    def _init_nodes(self, mine: Mine) -> None:
        for idx, i in enumerate(self.node_list):
            self.node_c[i] = [j for j in mine.graph[i] if j not in self.node_list]
            self.node_p[i] = [j for j in mine.graph_r[i] if j not in self.node_list]
            if len(self.node_c[i]) + len(self.node_p[i]) > 0:
                self.node_bd.append(idx)
                self.zb.append(single_qubit_pauli("z", idx, self.nqubits))

    def gen_Hs(self, mine: Mine) -> PauliOp:
        Hs = null_operator(self.nqubits)
        for sub_i, i in enumerate(self.node_list):
            for j in mine.graph[i]:
                try:
                    sub_j = self.node_list.index(j)
                    Hs += z_projector(1, sub_i, self.nqubits) @ z_projector(
                        0, sub_j, self.nqubits
                    )
                except:
                    continue
        return Hs

    def gen_Hp(self, mine: Mine) -> PauliOp:
        Hp = null_operator(self.nqubits)
        for sub_idx, idx in enumerate(self.node_list):
            Hp += float(mine.dat[mine.idx2cord[idx]]) * z_projector(
                1, sub_idx, self.nqubits
            )
        return Hp

    def _update_Hb(self, z_val: Dict) -> None:
        self.Hb = null_operator(self.nqubits)
        for sub_idx in self.node_bd:
            mine_idx = self.node_list[sub_idx]
            for j in self.node_c[mine_idx]:
                self.Hb += float(0.5 * (1 + z_val[j])) * z_projector(
                    1, sub_idx, self.nqubits
                )
            for j in self.node_p[mine_idx]:
                self.Hb += float(0.5 * (1 - z_val[j])) * z_projector(
                    0, sub_idx, self.nqubits
                )

    def gen_Hamiltonian(self, penalty: float, z_val: Dict) -> PauliOp:
        self._update_Hb(z_val)
        H_res = -self.Hp + penalty * (self.Hs + self.Hb)
        return H_res

    def solve(
        self, algorithm: MinimumEigensolver, penalty: Union[float, None], z_val: Dict
    ) -> "MiningProblemResult":
        if isinstance(algorithm, MinimumEigensolver):
            res = algorithm.compute_minimum_eigenvalue(
                self.gen_Hamiltonian(penalty, z_val), self.zb + [self.Hp, self.Hs]
            )
            self._ret = MiningProblemResult()
            self._ret.optimal_config, self._ret.optimal_config_prob = max(
                res.eigenstate.items(), key=lambda item: item[1]
            )
            # self._ret.ground_state = res.eigenstate
            for sub_idx, zi in zip(self.node_bd, res.aux_operator_eigenvalues[:-2]):
                z_val[self.node_list[sub_idx]] = zi
            self._ret.expected_profit = res.aux_operator_eigenvalues[-2]
            self._ret.expected_violation = res.aux_operator_eigenvalues[-1]
            return self._ret
        else:
            raise ValueError()

class SubMine(BaseMine):
    def __init__(self, mine: Mine, node_list: List) -> None:
        super.__init__()
        self.node_list = node_list
        self.nqubits = len(self.node_list)
        self.node_c = {}  # key: a node in this frag, value: child node in another frag
        self.node_p = {}  # key: a node in this frag, value: parent node in another frag
        self.node_bd = []  # submine idx of bd nodes
        self.zb = []  # pauli z op of bd nodes

        self._init_nodes(mine)
        self.Hp = self.gen_Hp(mine)  # profit
        self.Hs = self.gen_Hs(mine)  # smoothness, inner nodes
        self.Hb = null_operator(self.nqubits)  # boundary

    def _init_nodes(self, mine: Mine) -> None:
        for idx, i in enumerate(self.node_list):
            self.node_c[i] = [j for j in mine.graph[i] if j not in self.node_list]
            self.node_p[i] = [j for j in mine.graph_r[i] if j not in self.node_list]
            if len(self.node_c[i]) + len(self.node_p[i]) > 0:
                self.node_bd.append(idx)
                self.zb.append(single_qubit_pauli("z", idx, self.nqubits))

    def gen_Hs(self, mine: Mine) -> PauliOp:
        Hs = null_operator(self.nqubits)
        for sub_i, i in enumerate(self.node_list):
            for j in mine.graph[i]:
                try:
                    sub_j = self.node_list.index(j)
                    Hs += z_projector(1, sub_i, self.nqubits) @ z_projector(
                        0, sub_j, self.nqubits
                    )
                except:
                    continue
        return Hs

    def gen_Hp(self, mine: Mine) -> PauliOp:
        Hp = null_operator(self.nqubits)
        for sub_idx, idx in enumerate(self.node_list):
            Hp += float(mine.dat[mine.idx2cord[idx]]) * z_projector(
                1, sub_idx, self.nqubits
            )
        return Hp

    def _update_Hb(self, z_val: Dict) -> None:
        self.Hb = null_operator(self.nqubits)
        for sub_idx in self.node_bd:
            mine_idx = self.node_list[sub_idx]
            for j in self.node_c[mine_idx]:
                self.Hb += float(0.5 * (1 + z_val[j])) * z_projector(
                    1, sub_idx, self.nqubits
                )
            for j in self.node_p[mine_idx]:
                self.Hb += float(0.5 * (1 - z_val[j])) * z_projector(
                    0, sub_idx, self.nqubits
                )

    def gen_Hamiltonian(self, penalty: float, z_val: Dict) -> PauliOp:
        self._update_Hb(z_val)
        H_res = -self.Hp + penalty * (self.Hs + self.Hb)
        return H_res

    def solve(
        self, algorithm: MinimumEigensolver, penalty: Union[float, None], z_val: Dict
    ) -> "MiningProblemResult":
        if isinstance(algorithm, MinimumEigensolver):
            res = algorithm.compute_minimum_eigenvalue(
                self.gen_Hamiltonian(penalty, z_val), self.zb + [self.Hp, self.Hs]
            )
            self._ret = MiningProblemResult()
            self._ret.optimal_config, self._ret.optimal_config_prob = max(
                res.eigenstate.items(), key=lambda item: item[1]
            )
            # self._ret.ground_state = res.eigenstate
            for sub_idx, zi in zip(self.node_bd, res.aux_operator_eigenvalues[:-2]):
                z_val[self.node_list[sub_idx]] = zi
            self._ret.expected_profit = res.aux_operator_eigenvalues[-2]
            self._ret.expected_violation = res.aux_operator_eigenvalues[-1]
            return self._ret
        else:
            raise ValueError()

class MiningProblemResult(AlgorithmResult):
    def __init__(self) -> None:
        super().__init__()
        self._optimal_config = None
        self._optimal_config_prob = None
        self._ground_state = None
        self._expected_profit = None
        self._expected_violation = None

    @property
    def optimal_config(self) -> Optional[str]:
        return self._optimal_config

    @optimal_config.setter
    def optimal_config(self, optimal_config: str) -> None:
        self._optimal_config = optimal_config

    @property
    def optimal_config_prob(self) -> Optional[float]:
        return self._optimal_config_prob

    @optimal_config_prob.setter
    def optimal_config_prob(self, optimal_config_prob: float) -> None:
        self._optimal_config_prob = optimal_config_prob

    @property
    def ground_state(self) -> Optional[np.ndarray]:
        return self._ground_state

    @ground_state.setter
    def ground_state(self, ground_state: np.ndarray) -> None:
        self._ground_state = ground_state

    @property
    def expected_profit(self) -> Optional[float]:
        return self._expected_profit

    @expected_profit.setter
    def expected_profit(self, expected_profit: float) -> None:
        self._expected_profit = expected_profit

    @property
    def expected_violation(self) -> Optional[float]:
        return self._expected_violation

    @expected_violation.setter
    def expected_violation(self, expected_violation: float) -> None:
        self._expected_violation = expected_violation


"""
pretrain using qt, to get neural representation, then run kmeans on various
combinations of the resulting representations

this was forked initially from train.py, then modified
"""
import argparse
import datetime
import copy
import time
import numpy as np
import sklearn.cluster
import warnings
import torch
from torch import autograd

from proc_data import Dataset
from model.multiview_encoders import MultiviewEncoders
from metrics import cluster_metrics
import pretrain

warnings.filterwarnings(action='ignore', category=RuntimeWarning)

torch.manual_seed(0)
np.random.seed(0)

LSTM_LAYER = 1
LSTM_HIDDEN = 300
WORD_DROPOUT_RATE = 0.
DROPOUT_RATE = 0.
BATCH_SIZE = 32
LEARNING_RATE = 0.001

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def transform(data, model):
    model.eval()
    latent_zs = []
    n_batch = (len(data) + BATCH_SIZE - 1) // BATCH_SIZE
    for i in range(n_batch):
        data_batch = data[i*BATCH_SIZE:(i+1)*BATCH_SIZE]
        with autograd.no_grad():
            latent_z = model(data_batch, encoder='v1')
        latent_zs.append(latent_z.cpu().data.numpy())
    latent_zs = np.concatenate(latent_zs)
    return latent_zs

def calc_prec_rec_f1_acc(preds, golds):
    lgolds, lpreds = [], []
    for g, p in zip(golds, list(preds)):
        if g > 0:
            lgolds.append(g)
            lpreds.append(p)
    prec, rec, f1 = cluster_metrics.calc_prec_rec_f1(
        gnd_assignments=torch.LongTensor(lgolds).to(device),
        pred_assignments=torch.LongTensor(lpreds).to(device))
    acc = cluster_metrics.calc_ACC(
        torch.LongTensor(lpreds).to(device), torch.LongTensor(lgolds).to(device))
    return prec, rec, f1, acc

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data-path', type=str, default='./data/airlines_processed.csv')
    parser.add_argument('--glove-path', type=str, default='./data/glove.840B.300d.txt')
    parser.add_argument('--pre-epoch', type=int, default=5)
    parser.add_argument('--pt-batch', type=int, default=100)
    parser.add_argument('--scenarios', type=str, default='view1,view2,concatviews,wholeconv',
                        help='comma-separated, from [view1|view2|concatviews|wholeconv|mvsc]')
    parser.add_argument('--mvsc-no-unk', action='store_true',
                        help='only feed non-unk data to MVSC (to avoid oom)')

    parser.add_argument('--view1-col', type=str, default='view1')
    parser.add_argument('--view2-col', type=str, default='view2')
    parser.add_argument('--label-col', type=str, default='tag')
    args = parser.parse_args()

    print('loading dataset')
    dataset = Dataset(args.data_path, view1_col=args.view1_col, view2_col=args.view2_col,
                      label_col=args.label_col)
    n_cluster = len(dataset.id_to_label) - 1
    print("num of class = %d" % n_cluster)

    id_to_token, token_to_id = dataset.id_to_token, dataset.token_to_id
    vocab_size = len(dataset.token_to_id)
    print('vocab_size', vocab_size)

    # Load pre-trained GloVe vectors
    pretrained = {}
    word_emb_size = 0
    print('loading glove')
    for line in open(args.glove_path):
        parts = line.strip().split()
        if len(parts) % 100 != 1:
            continue
        word = parts[0]
        if word not in token_to_id:
            continue
        vector = [float(v) for v in parts[1:]]
        pretrained[word] = vector
        word_emb_size = len(vector)
    pretrained_list = []
    scale = np.sqrt(3.0 / word_emb_size)
    print('loading oov')
    for word in id_to_token:
        # apply lower() because all GloVe vectors are for lowercase words
        if word.lower() in pretrained:
            pretrained_list.append(np.array(pretrained[word.lower()]))
        else:
            random_vector = np.random.uniform(-scale, scale, [word_emb_size])
            pretrained_list.append(random_vector)

    model = MultiviewEncoders.from_embeddings(
        embeddings=torch.FloatTensor(pretrained_list),
        num_layers=LSTM_LAYER,
        embedding_size=word_emb_size,
        lstm_hidden_size=LSTM_HIDDEN,
        word_dropout=WORD_DROPOUT_RATE,
        dropout=DROPOUT_RATE,
        vocab_size=vocab_size
    )
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

    expressions = (model, optimizer)
    pre_acc, pre_state = 0., None
    pretrain_method = pretrain.pretrain_qt
    for epoch in range(1, args.pre_epoch + 1):
        model.train()
        perm_idx = np.random.permutation(dataset.trn_idx)
        trn_loss, _ = pretrain_method(dataset, perm_idx, expressions, train=True)
        model.eval()
        _, tst_acc = pretrain_method(dataset, dataset.tst_idx, expressions, train=False)
        if tst_acc > pre_acc:
            pre_state = copy.deepcopy(model.state_dict())
            pre_acc = tst_acc
        print(f'{datetime.datetime.now()} epoch {epoch}, train_loss={trn_loss:.4f} '
              f'test_acc={tst_acc:.4f}')

    if args.pre_epoch > 0:
        # load best state
        model.load_state_dict(pre_state)

        # deepcopy pretrained views into v1 and/or view2
        pretrain.after_pretrain_qt(model)

    kmeans = sklearn.cluster.KMeans(n_clusters=n_cluster, max_iter=300, verbose=0, random_state=0)

    golds = [dataset[idx][1] for idx in dataset.trn_idx]
    for rep in args.scenarios.split(','):
        if rep == 'view1':
            data = [dataset[idx][0][0] for idx in dataset.trn_idx]
            encoded = transform(data=data, model=model)
            preds = kmeans.fit_predict(encoded)
        elif rep == 'view2':
            data = [dataset[idx][0][1] for idx in dataset.trn_idx]
            encoded = []
            for conv in data:
                encoded_conv = transform(data=conv, model=model)
                encoded_conv = torch.from_numpy(encoded_conv)
                encoded_conv = encoded_conv.mean(dim=0)
                encoded.append(encoded_conv)
            encoded = torch.stack(encoded, dim=0)
            # print('encoded.size()', encoded.size())
            encoded = encoded.numpy()
            preds = kmeans.fit_predict(encoded)
        elif rep == 'concatviews':
            v1_data = [dataset[idx][0][0] for idx in dataset.trn_idx]
            v1_encoded = torch.from_numpy(transform(data=v1_data, model=model))

            v2_data = [dataset[idx][0][1] for idx in dataset.trn_idx]
            v2_encoded = []
            for conv in v2_data:
                encoded_conv = transform(data=conv, model=model)
                encoded_conv = torch.from_numpy(encoded_conv)
                encoded_conv = encoded_conv.mean(dim=0)
                v2_encoded.append(encoded_conv)
            v2_encoded = torch.stack(v2_encoded, dim=0)
            concatview = torch.cat([v1_encoded, v2_encoded], dim=-1)
            print('concatview.size()', concatview.size())
            encoded = concatview.numpy()
            preds = kmeans.fit_predict(encoded)
        elif rep == 'wholeconv':
            encoded = []
            for idx in dataset.trn_idx:
                v1 = dataset[idx][0][0]
                v2 = dataset[idx][0][1]
                conv = [v1] + v2
                encoded_conv = transform(data=conv, model=model)
                encoded_conv = torch.from_numpy(encoded_conv)
                encoded_conv = encoded_conv.mean(dim=0)
                encoded.append(encoded_conv)
            encoded = torch.stack(encoded, dim=0)
            print('encoded.size()', encoded.size())
            encoded = encoded.numpy()
            preds = kmeans.fit_predict(encoded)
        elif rep == 'mvsc':
            try:
                import multiview
            except Exception:
                print('please install https://github.com/mariceli3/multiview')
                return
            print('imported multiview ok')

            idx = dataset.trn_idx_no_unk if args.mvsc_no_unk else dataset.trn_idx
            v1_data = [dataset[idx][0][0] for idx in idx]
            v1_encoded = torch.from_numpy(transform(data=v1_data, model=model))

            v2_data = [dataset[idx][0][1] for idx in idx]
            v2_encoded = []
            for conv in v2_data:
                encoded_conv = transform(data=conv, model=model)
                encoded_conv = torch.from_numpy(encoded_conv)
                encoded_conv = encoded_conv.mean(dim=0)
                v2_encoded.append(encoded_conv)
            v2_encoded = torch.stack(v2_encoded, dim=0)

            mvsc = multiview.mvsc.MVSC(
                k=n_cluster
            )
            print('running mvsc', end='', flush=True)
            start = time.time()
            preds, eivalues, eivectors, sigmas = mvsc.fit_transform(
                [v1_encoded, v2_encoded], [False] * 2
            )
            print('...done')
            mvsc_time = time.time() - start
            print('time taken %.3f' % mvsc_time)
        else:
            raise Exception('unimplemented rep', rep)

        prec, rec, f1, acc = calc_prec_rec_f1_acc(preds, golds)
        print(f'{datetime.datetime.now()} {rep}: eval prec={prec:.4f} rec={rec:.4f} f1={f1:.4f} '
              f'acc={acc:.4f}')

if __name__ == '__main__':
    main()