import os
from pathlib import Path
import fnmatch
import subprocess
from collections import namedtuple
import tarfile
import logging

from ..handler import add_handler

import numpy as np
import xarray as xr
import rasterio as rio
from rasterio.crs import CRS
# from rasterio.warp import reproject, Resampling
from affine import Affine
import xml.etree.ElementTree as ET

try:
    import cv2
    OPENCV_INSTALLED = True
except:
    OPENCV_INSTALLED = False

logger = logging.getLogger(__name__)
logger = add_handler(logger)

def shift_objects(data,
                  solar_za,
                  solar_az,
                  sensor_za,
                  sensor_az,
                  h,
                  num_workers):

    """
    Shifts objects along x and y dimensions

    Args:
        data (DataArray): The data to shift.
        solar_za (DataArray): The solar zenith angle.
        solar_az (DataArray): The solar azimuth angle.
        sensor_za (DataArray): The sensor, or view, zenith angle.
        sensor_az (DataArray): The sensor, or view, azimuth angle.
        h (float): The object height.
        num_workers (Optional[int]): The number of dask workers.

    Returns:
        ``xarray.DataArray``
    """

    # Scale the angles to degrees
    sza = solar_za * 0.01
    sza.coords['band'] = [1]

    saa = solar_az * 0.01
    saa.coords['band'] = [1]

    vza = sensor_za * 0.01
    vza.coords['band'] = [1]

    vaa = sensor_az * 0.01
    vaa.coords['band'] = [1]

    # Convert to radians
    rad_sza = np.deg2rad(sza)
    rad_saa = np.deg2rad(saa)
    rad_vza = np.deg2rad(vza)
    rad_vaa = np.deg2rad(vaa)

    apparent_solar_az = np.pi + np.arctan((np.sin(rad_saa) * np.tan(rad_sza) - np.sin(rad_vaa) * np.tan(rad_vza)) /
                                          (np.cos(rad_saa) * np.tan(rad_sza) - np.cos(rad_vaa) * np.tan(rad_vza)))

    # Maximum horizontal distance
    d = (h**2 * ((np.sin(rad_saa) * np.tan(rad_sza) - np.sin(rad_vaa) * np.tan(rad_vza))**2 +
                 (np.cos(rad_saa) * np.tan(rad_sza) - np.cos(rad_vaa) * np.tan(rad_vza))**2))**0.5

    # Convert the polar angle to cartesian offsets
    x = int((np.cos(apparent_solar_az) * d).max(skipna=True).data.compute(num_workers=num_workers))
    y = int((np.sin(apparent_solar_az) * d).max(skipna=True).data.compute(num_workers=num_workers))

    return data.shift(shifts={'x': x, 'y': y}, fill_value=0)

def estimate_cloud_shadows(data,
                           clouds,
                           solar_za,
                           solar_az,
                           sensor_za,
                           sensor_az,
                           heights=None,
                           num_workers=1):

    """
    Estimates shadows from a cloud mask and adds to the existing mask

    Args:
        data (DataArray): The wavelengths, scaled 0-1.
        clouds (DataArray): The cloud mask, where clouds=1 and clear sky=0.
        solar_za (DataArray): The solar zenith angle.
        solar_az (DataArray): The solar azimuth angle.
        sensor_za (DataArray): The sensor, or view, zenith angle.
        sensor_az (DataArray): The sensor, or view, azimuth angle.
        heights (Optional[list]): The cloud heights, in kilometers.
        num_workers (Optional[int]): The number of dask workers.

    Returns:
        ``xarray.DataArray``

    References:

        For the angle offset calculations, see :cite:`fisher_2014`.
        For the shadow test, see :cite:`sun_etal_2018`.
    """

    attrs = data.attrs.copy()

    if not heights:
        heights = list(range(200, 1400, 200))

    shadows = None

    for h in heights:

        potential_shadows = shift_objects(clouds,
                                          solar_za,
                                          solar_az,
                                          sensor_za,
                                          sensor_az,
                                          h,
                                          num_workers)

        if not isinstance(shadows, xr.DataArray):
            shadows = xr.where((data.sel(band='nir') < 0.25) & (data.sel(band='swir1') < 0.11) & (potential_shadows.sel(band='mask') == 1), 1, 0)
        else:
            shadows = xr.where(((data.sel(band='nir') < 0.25) & (data.sel(band='swir1') < 0.11) & (potential_shadows.sel(band='mask') == 1)) | shadows.sel(band='mask') == 1, 1, 0)

        shadows = shadows.expand_dims(dim='band')

    # Add the shadows to the cloud mask
    data = xr.where(clouds.sel(band='mask') == 1, 1, xr.where(shadows.sel(band='mask') == 1, 2, 0))
    data = data.expand_dims(dim='band')
    data.attrs = attrs

    return data

def scattering_angle(cos_sza, cos_vza, sin_sza, sin_vza, cos_raa):

    """
    Calculates the scattering angle

    Args:
        cos_sza (DataArray): The cosine of the solar zenith angle.
        cos_vza (DataArray): The cosine of the view zenith angle.
        sin_sza (DataArray): The sine of the solar zenith angle.
        sin_vza (DataArray): The sine of the view zenith angle.
        cos_raa (DataArray): The cosine of the relative azimuth angle.

    Equation:

        .. math::

            \Theta = scattering angle

            \theta_0 = solar zenith angle

            \theta_S = sensor zenith angle

            \zeta = relative azimuth angle

            \Theta_s = \arccos{- \cos{\theta_0} \cos{\theta_S} - \sin{\theta_0} \sin{\theta_S} \cos{\zeta}}

    References:
        scattering angle = the angle between the direction of incident and scattered radiation
        Liu, CH and Liu GR (2009) AEROSOL OPTICAL DEPTH RETRIEVAL FOR SPOT HRV IMAGES, Journal of Marine Science and Technology
        http://stcorp.github.io/harp/doc/html/algorithms/derivations/scattering_angle.html

    Returns:
        Scattering angle (in radians) as an ``xarray.DataArray``
    """

    scattering_angle = xr.ufuncs.arccos(-cos_sza * cos_vza - sin_sza * sin_vza * cos_raa)

    return xr.ufuncs.cos(scattering_angle) ** 2

def relative_azimuth(saa, vaa):

    """
    Calculates the relative azimuth angle

    Args:
        saa (DataArray): The solar azimuth angle (in degrees).
        vaa (DataArray): The view azimuth angle (in degrees).

    Reference:
        http://stcorp.github.io/harp/doc/html/algorithms/derivations/relative_azimuth_angle.html

    Returns:
        Relative azimuth (in degrees) as an ``xarray.DataArray``
    """

    # Relative azimuth (in radians)
    raa = xr.ufuncs.deg2rad(saa - vaa)

    # Create masks
    raa_plus = xr.where(raa >= 2.0*np.pi, 1, 0)
    raa_minus = xr.where(raa < 0, 1, 0)

    # raa = xr.where(raa_plus == 1, raa + (2.0*np.pi), raa)
    # raa = xr.where(raa_minus == 1, raa - (2.0*np.pi), raa)

    raa = xr.where(raa_plus == 1, raa - (2.0 * np.pi), raa)
    raa = xr.where(raa_minus == 1, raa + (2.0 * np.pi), raa)

    return xr.ufuncs.fabs(xr.ufuncs.rad2deg(raa))

def get_sentinel_sensor(metadata):

    # Parse the XML file
    tree = ET.parse(metadata)
    root = tree.getroot()

    for child in root:

        if 'general_info' in child.tag[-14:].lower():
            general_info = child

    for ginfo in general_info:

        if ginfo.tag == 'TILE_ID':
            file_name = ginfo.text

    return file_name[:3].lower()

def parse_sentinel_angles(metadata, proc_angles, nodata):

    """
    Gets the Sentinel-2 solar angles from metadata

    Args:
        metadata (str): The metadata file.
        proc_angles (str): The angles to parse. Choices are ['solar', 'view'].
        nodata (int or float): The 'no data' value.

    Returns:
        zenith and azimuth angles as a ``tuple`` of 2d ``numpy`` arrays
    """

    if proc_angles == 'view':

        zenith_values = np.zeros((13, 23, 23), dtype='float64') + nodata
        azimuth_values = np.zeros((13, 23, 23), dtype='float64') + nodata

    else:

        zenith_values = np.zeros((23, 23), dtype='float64') + nodata
        azimuth_values = np.zeros((23, 23), dtype='float64') + nodata

    view_tag = 'Sun_Angles_Grid' if proc_angles == 'solar' else 'Viewing_Incidence_Angles_Grids'

    # Parse the XML file
    tree = ET.parse(metadata)
    root = tree.getroot()

    # Find the angles
    for child in root:

        if child.tag.split('}')[-1] == 'Geometric_Info':
            geoinfo = child
            break

    for segment in geoinfo:

        if segment.tag == 'Tile_Angles':
            angles = segment

    for angle in angles:

        if angle.tag == view_tag:

            if proc_angles == 'view':
                band_id = int(angle.attrib['bandId'])

            for bset in angle:

                if bset.tag == 'Zenith':
                    zenith = bset
                if bset.tag == 'Azimuth':
                    azimuth = bset

            for field in zenith:

                if field.tag == 'Values_List':
                    zvallist = field

            for field in azimuth:

                if field.tag == 'Values_List':
                    avallist = field

            for rindex in range(len(zvallist)):

                zvalrow = zvallist[rindex]
                avalrow = avallist[rindex]
                zvalues = zvalrow.text.split(' ')
                avalues = avalrow.text.split(' ')
                values = list(zip(zvalues, avalues))

                for cindex in range(len(values)):

                    if (values[cindex][0].lower() != 'nan') and (values[cindex][1].lower() != 'nan'):

                        ze = float(values[cindex][0])
                        az = float(values[cindex][1])

                        if proc_angles == 'view':

                            zenith_values[band_id, rindex, cindex] = ze
                            azimuth_values[band_id, rindex, cindex] = az

                        else:

                            zenith_values[rindex, cindex] = ze
                            azimuth_values[rindex, cindex] = az

    return zenith_values, azimuth_values

def sentinel_pixel_angles(metadata,
                          ref_file,
                          outdir='.',
                          nodata=-32768,
                          overwrite=False,
                          verbose=0):

    """
    Generates Sentinel pixel angle files

    Args:
        metadata (str): The metadata file.
        ref_file (str): A reference image to use for geo-information.
        outdir (Optional[str])): The output directory to save the angle files to.
        nodata (Optional[int or float]): The 'no data' value.
        overwrite (Optional[bool]): Whether to overwrite existing angle files.
        verbose (Optional[int]): The verbosity level.

    References:
        https://www.sentinel-hub.com/faq/how-can-i-access-meta-data-information-sentinel-2-l2a
        https://github.com/marujore/sentinel_angle_bands/blob/master/sentinel2_angle_bands.py

    Returns:
        zenith and azimuth angles as a ``namedtuple`` of angle file names
    """

    if not OPENCV_INSTALLED:
        logger.exception('OpenCV must be installed.')

    AngleInfo = namedtuple('AngleInfo', 'vza vaa sza saa sensor')

    sza, saa = parse_sentinel_angles(metadata, 'solar', nodata)
    vza, vaa = parse_sentinel_angles(metadata, 'view', nodata)

    sensor_name = get_sentinel_sensor(metadata)

    with rio.open(ref_file) as src:

        profile = src.profile.copy()

        ref_height = src.height
        ref_width = src.width
        ref_extent = src.bounds

        profile.update(transform=Affine(src.res[0], 0.0, ref_extent.left, 0.0, -src.res[1], ref_extent.top),
                       height=ref_height,
                       width=ref_width,
                       nodata=-32768,
                       dtype='int16',
                       count=1,
                       driver='GTiff',
                       tiled=True,
                       compress='lzw')

    ref_base = '_'.join(os.path.basename(ref_file).split('_')[:-1])

    opath = Path(outdir)

    opath.mkdir(parents=True, exist_ok=True)

    # Set output angle file names.
    sensor_azimuth_file = opath.joinpath(ref_base + '_sensor_azimuth.tif').as_posix()
    sensor_zenith_file = opath.joinpath(ref_base + '_sensor_zenith.tif').as_posix()
    solar_azimuth_file = opath.joinpath(ref_base + '_solar_azimuth.tif').as_posix()
    solar_zenith_file = opath.joinpath(ref_base + '_solar_zenith.tif').as_posix()

    for angle_array, angle_file in zip([vaa,
                                        vza,
                                        saa,
                                        sza],
                                       [sensor_azimuth_file,
                                        sensor_zenith_file,
                                        solar_azimuth_file,
                                        solar_zenith_file]):

        pfile = Path(angle_file)

        if overwrite:

            if pfile.is_file():
                pfile.unlink()

        if not pfile.is_file():

            # TODO: write data for each band?
            if len(angle_array.shape) > 2:
                angle_array = angle_array.mean(axis=0)

            with rio.open(angle_file, mode='w', **profile) as dst:

                if verbose > 0:
                    logger.info('  Writing {} to file ...'.format(angle_file))

                # Resample and scale
                angle_array_resamp = np.int16(cv2.resize(angle_array,
                                                         (0, 0),
                                                         fy=ref_height / angle_array.shape[0],
                                                         fx=ref_width / angle_array.shape[1],
                                                         interpolation=cv2.INTER_LINEAR) / 0.01)

                dst.write(angle_array_resamp, indexes=1)

    return AngleInfo(vaa=str(sensor_azimuth_file),
                     vza=str(sensor_zenith_file),
                     saa=str(solar_azimuth_file),
                     sza=str(solar_zenith_file),
                     sensor=sensor_name)

# Potentially useful for angle creation
# https://github.com/gee-community/gee_tools/blob/master/geetools/algorithms.py

# def slope_between(a, b):
#     return (a[1] - b[1]) / (a[0] - b[0])
#
#
# @nb.jit
# def _calc_sensor_angles(data,
#                         zenith_angles,
#                         azimuth_angles,
#                         yvalues,
#                         xvalues,
#                         celly,
#                         cellx,
#                         satellite_height,
#                         nodata,
#                         acquisition_date):
#
#     """
#     Calculates sensor zenith and azimuth angles
#     """
#
#     slope = slope_between(np.array([data.gw.meta.right + ((data.gw.ncols/2.0)*data.gw.cellx), data.gw.meta.top]),
#                           np.array([data.gw.meta.left, data.gw.meta.top - ((data.gw.nrows / 2.0) * data.gw.celly)]))
#
#     slope_perc = -1.0 / slope
#
#     view_az = (math.pi / 2.0) - math.arctan(slope_perc)
#
#     for i in range(0, yvalues.shape[0]):
#
#         for j in range(0, xvalues.shape[0]):
#
#             if data_band[i, j] != nodata:
#
#                 # TODO: calculate satellite drift angle
#                 dist_from_nadir = None
#
#                 # Calculate the distance from the current location to the satellite
#                 dist_to_satellite = np.hypot(satellite_height, dist_from_nadir)
#
#                 # Calculate the view angle
#
#                 zenith_angles[i, j]
#
#                 # Solar zenith angle = 90 - elevation angle scaled to integer range
#                 zenith_angles[i, j] = (90.0 - get_altitude_fast(xvalues[j], yvalues[i], acquisition_date)) / 0.01
#
#                 # Solar azimuth angle
#                 azimuth_angles[i, j] = float(get_azimuth_fast(xvalues[j], yvalues[i], acquisition_date)) / 0.01
#
#     return zenith_angles, azimuth_angles

# @nb.jit
# def _calc_solar_angles(data_band, zenith_angles, azimuth_angles, yvalues, xvalues, nodata, acquisition_date):
#
#     """
#     Calculates solar zenith and azimuth angles
#     """
#
#     for i in range(0, yvalues):
#
#         for j in range(0, xvalues):
#
#             if data_band[i, j] != nodata:
#
#                 # Solar zenith angle = 90 - elevation angle scaled to integer range
#                 zenith_angles[i, j] = (90.0 - get_altitude_fast(xvalues[j], yvalues[i], acquisition_date)) / 0.01
#
#                 # Solar azimuth angle
#                 azimuth_angles[i, j] = float(get_azimuth_fast(xvalues[j], yvalues[i], acquisition_date)) / 0.01
#
#     return zenith_angles, azimuth_angles

# def pixel_angles(data, band, nodata, meta):
#
#     """
#     Generates pixel zenith and azimuth angles
#
#     Args:
#         data (Xarray): The data with coordinate and transform attributes.
#         band (int or str): The ``data`` band to use for masking.
#         nodata (int or float): The 'no data' value in ``data``.
#         meta (namedtuple): The metadata file. Should have image acquisition year, month, day and hour attributes.
#     """
#
#     acquisition_date = dtime(meta.year, meta.month, meta.day, meta.hour, 0, 0, 0, tzinfo=datetime.timezone.utc)
#
#     yvalues = data.y.values
#     xvalues = data.x.values
#
#     data_band = data.sel(band=band).data.compute()
#     sze = np.zeros((data.gw.nrows, data.gw.ncols), dtype='int16') - 32768
#     saa = np.zeros((data.gw.nrows, data.gw.ncols), dtype='int16') - 32768
#
#     sze, saa = _calc_solar_angles(data_band, sze, saa, yvalues, xvalues, nodata, acquisition_date)
#
#     sze_attrs = data.attrs.copy()
#     saa_attrs = data.attrs.copy()
#
#     sze_attrs['values'] = 'Solar zenith angle'
#     sze_attrs['scale_factor'] = 0.01
#
#     saa_attrs['values'] = 'Solar azimuth angle'
#     sze_attrs['scale_factor'] = 0.01
#
#     szex = xr.DataArray(data=da.from_array(sze[np.newaxis, :, :],
#                                            chunks=(1, data.gw.row_chunks, data.gw.col_chunks)),
#                         coords={'band': 'sze',
#                                 'y': data.y,
#                                 'x': data.x},
#                         dims=('band', 'y', 'x'),
#                         attrs=sze_attrs)
#
#     saax = xr.DataArray(data=da.from_array(saa[np.newaxis, :, :],
#                                            chunks=(1, data.gw.row_chunks, data.gw.col_chunks)),
#                         coords={'band': 'saa',
#                                 'y': data.y,
#                                 'x': data.x},
#                         dims=('band', 'y', 'x'),
#                         attrs=saa_attrs)
#
#     return szex, saax

def landsat_pixel_angles(angles_file,
                         ref_file,
                         out_dir,
                         sensor,
                         l57_angles_path=None,
                         l8_angles_path=None,
                         subsample=1,
                         resampling='bilinear',
                         num_threads=1,
                         verbose=0):

    """
    Generates Landsat pixel angle files

    Args:
        angles_file (str): The angles file.
        ref_file (str): A reference file.
        out_dir (str): The output directory.
        sensor (str): The sensor.
        l57_angles_path (str): The path to the Landsat 5 and 7 angles bin.
        l8_angles_path (str): The path to the Landsat 8 angles bin.
        subsample (Optional[int]): The sub-sample factor when calculating the angles.
        resampling (Optional[str]): The resampling method if ``filename`` is a ``list``.
            Choices are ['average', 'bilinear', 'cubic', 'cubic_spline', 'gauss', 'lanczos', 'max', 'med', 'min', 'mode', 'nearest'].
        num_threads (Optional[int]): The number of threads to pass to ``rasterio.warp.reproject``.
        verbose (Optional[int]): The verbosity level.

    Returns:
        zenith and azimuth angles as a ``namedtuple`` of angle file names
    """

    if not l57_angles_path:

        gw_bin = os.path.realpath(os.path.dirname(__file__))

        gw_out = os.path.realpath(Path(gw_bin).joinpath('../bin').as_posix())
        gw_tar = os.path.realpath(Path(gw_bin).joinpath('../bin/ESPA.tar.gz').as_posix())

        if not Path(gw_bin).joinpath('../bin/ESPA').is_dir():

            with tarfile.open(gw_tar, mode='r:gz') as tf:
                tf.extractall(gw_out)

        l57_angles_path = Path(gw_out).joinpath('ESPA/landsat_angles').as_posix()
        l8_angles_path = Path(gw_out).joinpath('ESPA/l8_angles').as_posix()

    AngleInfo = namedtuple('AngleInfo', 'vza vaa sza saa')

    # Setup the angles name.
    # example file = LE07_L1TP_225098_20160911_20161008_01_T1_sr_band1.tif

    with rio.open(ref_file) as src:

        ref_res = src.res
        ref_height = src.height
        ref_width = src.width
        ref_extent = src.bounds

    ref_base = '_'.join(os.path.basename(ref_file).split('_')[:-1])

    opath = Path(out_dir)

    opath.mkdir(parents=True, exist_ok=True)

    # Set output angle file names.
    sensor_azimuth_file = opath.joinpath(ref_base + '_sensor_azimuth.tif').as_posix()
    sensor_zenith_file = opath.joinpath(ref_base + '_sensor_zenith.tif').as_posix()
    solar_azimuth_file = opath.joinpath(ref_base + '_solar_azimuth.tif').as_posix()
    solar_zenith_file = opath.joinpath(ref_base + '_solar_zenith.tif').as_posix()

    if not Path(sensor_azimuth_file).is_file():

        # Setup the command.
        if sensor.lower() in ['l5', 'l7']:

            angle_command = '{PATH} {META} -s {SUBSAMP:d} -b 1'.format(PATH=str(Path(l57_angles_path).joinpath('landsat_angles')),
                                                                       META=angles_file,
                                                                       SUBSAMP=subsample)

            # 1=zenith, 2=azimuth
            out_order = dict(azimuth=2, zenith=1)
            # out_order = [2, 1, 2, 1]

        else:

            angle_command = '{PATH} {META} BOTH {SUBSAMP:d} -f -32768 -b 4'.format(PATH=str(Path(l8_angles_path).joinpath('l8_angles')),
                                                                                META=angles_file,
                                                                                SUBSAMP=subsample)

            # 1=azimuth, 2=zenith
            out_order = dict(azimuth=1, zenith=2)
            # out_order = [1, 2, 1, 2]

        os.chdir(out_dir)

        if verbose > 0:
            logger.info('  Generating pixel angles ...')

        # Create the angle files.
        subprocess.call(angle_command, shell=True)

        # Get angle data from 1 band.
        sensor_angles = fnmatch.filter(os.listdir(out_dir), '*sensor_B04.img')[0]
        solar_angles = fnmatch.filter(os.listdir(out_dir), '*solar_B04.img')[0]

        sensor_angles_fn_in = opath.joinpath(sensor_angles).as_posix()
        solar_angles_fn_in = opath.joinpath(solar_angles).as_posix()

        # Convert the data
        for in_angle, out_angle, band_pos in zip([sensor_angles_fn_in,
                                                  sensor_angles_fn_in,
                                                  solar_angles_fn_in,
                                                  solar_angles_fn_in],
                                                 [sensor_azimuth_file,
                                                  sensor_zenith_file,
                                                  solar_azimuth_file,
                                                  solar_zenith_file],
                                                 ['azimuth',
                                                  'zenith',
                                                  'azimuth',
                                                  'zenith']):

            new_res = subsample*ref_res[0]

            # Update the .hdr file
            with open(in_angle + '.hdr', mode='r') as txt:

                lines = txt.readlines()

                for lidx, line in enumerate(lines):
                    if line.startswith('map info'):
                        lines[lidx] = line.replace('30.000, 30.000', f'{new_res:.3f}, {new_res:.3f}')

            Path(in_angle + '.hdr').unlink()

            with open(in_angle + '.hdr', mode='w') as txt:
                txt.writelines(lines)

            with rio.open(in_angle) as src:

                profile = src.profile.copy()

                epsg = src.crs.to_epsg()

                # Adjust Landsat images in the Southern hemisphere
                if str(epsg).startswith('326') and (ref_extent.top < 0):

                    transform = Affine(new_res, 0.0, ref_extent.left, 0.0, -new_res, ref_extent.top+10_000_000.0)
                    crs = CRS.from_epsg(f'327{str(epsg)[3:]}')

                else:

                    transform = Affine(new_res, 0.0, ref_extent.left, 0.0, -new_res, ref_extent.top)
                    crs = src.crs

                profile.update(transform=transform,
                               crs=crs,
                               height=src.height,
                               width=src.width,
                               nodata=-32768,
                               dtype='int16',
                               count=1,
                               driver='GTiff',
                               tiled=True,
                               compress='lzw')

                # src_band = rio.Band(src, out_order[band_pos], 'int16', (src.height, src.width))

                with rio.open(out_angle, mode='w', **profile) as dst:

                    dst.write(src.read(out_order[band_pos]),
                              indexes=1)

                    # dst_band = rio.Band(dst, 1, 'int16', (dst.height, dst.width))
                    #
                    # reproject(src_band,
                    #           destination=dst_band,
                    #           resampling=getattr(Resampling, resampling),
                    #           num_threads=num_threads)

    os.chdir(os.path.expanduser('~'))

    return AngleInfo(vaa=str(sensor_azimuth_file),
                     vza=str(sensor_zenith_file),
                     saa=str(solar_azimuth_file),
                     sza=str(solar_zenith_file))


import logging_setup
import logging
import numpy as np
import openfermion as of
import os
import qiskit as qk
import qiskit.opflow as qk_opflow
import qiskit.quantum_info as qk_qi
import uccsd_evolution
import scipy.linalg as spla
import warnings

try:
    from hubbard_bqskit import BQSKit_Hubbard_Optimizer
except ImportError:
    pass

__all__ = [
    'EnergyObjective',
    'hamiltonian_matrix',
    'hamiltonian_qiskit',
    'small_model',
    'medium_model',
    'clear_circuit_cache',
    'get_cached_circuit',
]

logger = logging.getLogger('hubbard')

# Allow caching of the trotterized opflow, post-BQSKit, for noise studies
_cached_circuit = None

def clear_circuit_cache():
    global _cached_circuit
    _cached_circuit = None

def get_cached_circuit():
   return _cached_circuit

class EnergyObjective:
    def __init__(self, hamiltonian, n_electrons_up, n_electrons_down,
                 trotter_steps=2, noise_model=None, shots=-1,
                 run_bqskit=False, save_evals=None):
        """\
        Create an energy estimater for the given Hamiltonian

        Args:
            hamiltonian(opflow): Hamiltonian operator
            n_electrons_up(int): number of spin-up electrons in the physical system
            n_electrons_down(int): number of spin-down electrons in the physical system
            trotter_steps(int): number of Trotter time slices for the evolution
            noise_model(NoiseModel): Qiskit noise model to apply
            shots(int): number of shots to sample and average over
            run_bqskit(bool): whether to run the bqskit stack on the evolution operator
            save_evals(str): file name to store evaluations or None
        """

        self._hamiltonian = hamiltonian
        self._n_qubits         = hamiltonian.num_qubits
        self._n_electrons      = n_electrons_up + n_electrons_down

        try:
            self._fermion_transform = hamiltonian._fermion_transform
        except AttributeError as a:
            self._fermion_transform = 'jordan-wigner'

      # Create initial state and add electrons by setting qubits to |1> (i.e., occupied)
        reg = qk.QuantumCircuit(self._n_qubits)
        if self._fermion_transform == 'bravyi-kitaev':
            if self._n_electrons:
              # fill out the mapping of electrons
                m = [0]*self._n_qubits
                for i in range(n_electrons_up):
                    m[i*2] = 1
                for i in range(n_electrons_down):
                    m[i*2+1] = 1

                for i in range(self._n_qubits):
                    if i % 2:                              # odd
                        if sum(m[:i+1]) % 2: reg.x(i)
                    elif m[i]:                             # even
                        reg.x(i)

        elif self._fermion_transform == 'jordan-wigner':
            for i in range(n_electrons_up):
                reg.x(i*2)
            for i in range(n_electrons_down):
                reg.x(i*2+1)

        self._state_in = qk_opflow.CircuitStateFn(reg)

      # Create an observable from the Hamiltonian
        self._meas_op = qk_opflow.StateFn(self._hamiltonian, is_measurement=True)

      # Number of Trotter steps to use in the evolution operator (see __call__)
        self._trotter_steps = trotter_steps

      # Create the simulator
        self._expectation = qk_opflow.PauliExpectation()
        if shots <= 0 and noise_model is None:
            self._simulator = None
            self._meas_components = None
        else:
            if noise_model is None:
                backend = qk.Aer.get_backend('qasm_simulator')
            else:
              # nominally, options should pass through kwargs of get_backend, however,
              # this does not appear to work for Aer, so set the noise_model option
              # explicitly on the retrieved backend
                if type(noise_model) == str:
                  # use an existing, named, IBM backend from the qiskit test suite to
                  # create a realistic noise model; if 'realistic' default to Montreal
                    if noise_model.lower() == 'realistic':
                        noise_model = 'Montreal'

                    import qiskit.test.mock as qk_mock
                    import qiskit.providers.aer as qk_aer_provides

                    fake_backend = 'Fake'+noise_model[0].upper()+noise_model[1:]
                    fake_device  = getattr(qk_mock, fake_backend)()

                    backend = qk_aer_provides.AerSimulator.from_backend(fake_device)
                else:
                    backend = qk.Aer.get_backend('aer_simulator', noise_model=noise_model)
                    backend.set_options(noise_model=noise_model)

                if shots <= 0:
                  # if not simulating sampling, use the AerPauliExpectation, which computes
                  # the expectation value given the noise (effectively "infinite" sampling);
                  # it also passes a special "instruction" to the sampler to ignore shots
                     self._expectation = qk_opflow.AerPauliExpectation()
                     shots = 2**20          # i.e. large enough not to contribute

            self._simulator = qk_opflow.CircuitSampler(backend=backend)
            self._simulator.quantum_instance.run_config.shots = shots

          # split measurement components to prevent fake coherent errors
            primitive = self._meas_op.primitive
            try:
                while 1: primitive = primitive.primitive
            except AttributeError:
                pass

            self._meas_components = list()
            for ops, coeff in primitive.to_list():
                self._meas_components.append(qk_opflow.StateFn(
                    qk_opflow.PauliOp(qk.quantum_info.Pauli(ops), coeff), is_measurement=True))

      # Flag to toggle running the BQSKit optimizer on the evolution operator (see __call__)
        self.bqskit_opt = BQSKit_Hubbard_Optimizer(run_bqskit=='full') if run_bqskit else None

      # File name to store evaluations, if requested
        if save_evals:
            self._save_evals = type(save_evals) == str and save_evals or 'pointlog.txt'
            try:
                os.remove(self._save_evals)
            except Exception:
                pass
        else:
            self._save_evals = None

    def npar(self):
        """\
        Number of independent parameters for the optimizer

        Returns:
           npar(int): number of parameters for the optimizer
        """

        return uccsd_evolution.singlet_paramsize(self._n_qubits, self._n_electrons)

    def generate_evolution_op(self, packed_amplitudes):
        """\
        Construct the evolution operator

        Returns:
            trotterized_ev_op (opflow): (trotterized, optimized) evolution operator
        """

      # Build the state preparation evolution operator
        if self._fermion_transform == 'bravyi-kitaev':
            def bk_with_qubits(fop):
                return of.transforms.bravyi_kitaev(fop, self._n_qubits)
            fermion_transform = bk_with_qubits
        elif self._fermion_transform == 'jordan-wigner':
            fermion_transform = of.transforms.jordan_wigner

        evolution_op = uccsd_evolution.singlet_evolution(
                           packed_amplitudes, self._n_qubits, self._n_electrons,
                           fermion_transform=fermion_transform)

      # Trotterize the evolution operator flow to be able to construct a circuit (the
      # choice of 2 slices was empirically determined; it may not fit all cases)
        if 0 < self._trotter_steps:
            num_time_slices = self._trotter_steps
            trotterized_ev_op = qk_opflow.PauliTrotterEvolution(
                trotter_mode='trotter', reps=num_time_slices).convert(evolution_op)

          # Run bqskit circuit optimizers as requested (only works on the trotterized
          # evolution operator as the normal time evolution is not unitary)
            if self.bqskit_opt is not None:
                trotterized_ev_op = self.bqskit_opt.optimize_evolution(trotterized_ev_op)

        else:
            trotterized_ev_op = evolution_op

        return trotterized_ev_op

    def generate_circuit(self, packed_amplitudes):
        """\
        Construct the circuit for the current parameters

        For the given packed_amplitudes, return the circuit to execute if this
        was a step in a VQE algorithm. The measurements are left out, because
        calculating a single Hamilitonian requires measurement Pauli-strings,
        many of which are independent.

        Args:
            packed_amplitudes(ndarray): compact array storing the unique single
                and double excitation amplitudes for a singlet UCCSD opflow.
                The ordering lists unique single excitations before double
                excitations

        Returns:
            circuit (QuantumCircuit): circuit for the current parameters
        """

      # Build the state preparation evolution operator
        trotterized_ev_op = self.generate_evolution_op(packed_amplitudes)

      # Combine with initializer and evolution
        expect_op = self._expectation.convert(
                            trotterized_ev_op @ self._state_in
                    )

      # Convert to QuantumCircuit
        circuit = (trotterized_ev_op @ self._state_in).to_circuit()

        return circuit

    def __call__(self, packed_amplitudes, use_cached_circuit=None):
        """\
        Calculate the energy expectation for the given parameters

        Args:
            packed_amplitudes(ndarray): compact array storing the unique single
                and double excitation amplitudes for a singlet UCCSD opflow.
                The ordering lists unique single excitations before double
                excitations
            use_cached_circuit(bool): use an existing cached circuit, or cache
                the currently calculated circuit

        Returns:
            energy(float): energy estimate
        """

      # Build the state preparation evolution operator
        global _cached_circuit
        if use_cached_circuit and _cached_circuit is not None:
            trotterized_ev_op = _cached_circuit
        elif isinstance(use_cached_circuit, qk_opflow.operator_base.OperatorBase):
            trotterized_ev_op = use_cached_circuit
        else:
            trotterized_ev_op = self.generate_evolution_op(packed_amplitudes)
            if use_cached_circuit:
                _cached_circuit = trotterized_ev_op

      # Run full simulation. If there are no errors, take a short cut and evaluate
      # the hamiltonian directly. Otherwise, to prevent unrealistic coherent errors,
      # calculate the energy from its components
        if self._simulator is None:
          # exact calculation
            expect_op = self._expectation.convert(
                            self._meas_op @ trotterized_ev_op @ self._state_in
                        )
            energy = np.real(expect_op.eval())

        else:
          # sampled calculation from components

          # Note, sampling of the full hamiltonian would look like:
          #     sampled_op = self._simulator.convert(
          #                      self._expectation.convert(
          #                          self._meas_op @ trotterized_ev_op @ self._state_in
          #                      )
          #                  )
          #     energy = np.real(sampled_op.eval())

          # use measurement components to prevent fake coherent errors
            energy = 0.
            for meas_op in self._meas_components:
                sampled_op = self._simulator.convert(
                                 self._expectation.convert(
                                     meas_op @ trotterized_ev_op @ self._state_in
                                 )
                             )
                energy += np.real(sampled_op.eval())

        logger.info('objective: %.5f @ %s', energy, packed_amplitudes)

        if self._save_evals:
          # store parameter values and energy in log file
            f = open(self._save_evals, "a+")
            for ii in range(len(packed_amplitudes)):
                f.write("%f  " % (packed_amplitudes[ii]))
            f.write("%f \n" % (energy))
            f.close()

        return energy

def _to_qiskit(of_qop, n_qubits):
    """Convert OpenFermion QubitOperators to Qiskit equivalent"""

    opflow = list()
    for paulis, coeff in sorted(of_qop.terms.items()):
        ops = ['I']*n_qubits
        for term in paulis:
            ops[term[0]] = term[1]

        ops.reverse()

        opflow1 = coeff*getattr(qk_opflow, ops[0])
        for i in range(1, n_qubits):
            opflow1 ^= getattr(qk_opflow, ops[i])
        opflow.append(opflow1)

    return sum(opflow)

def _hubbard_qubit(x_dimension, y_dimension, tunneling, coulomb,
        chemical_potential = 0.00, magnetic_field = 0.0, periodic = True, spinless = False,
        fermion_transform='bravyi-kitaev'):
    """Create Fermi-Hubbard model with OpenFermion"""

    _fermion_transform = fermion_transform.lower()
    known_transforms = ['jordan-wigner', 'bravyi-kitaev']
    if not _fermion_transform in known_transforms:
        raise ValueError("unknown transform '%s'" % fermion_transform)

    # Hubbard Hamiltonian expressed in FermionOperators, i.e. creation and annihilation
    # operators. Each FermionOperator consists of the site it operates on (expressed as
    # an "index") and whether it raises (1) or lowers (0; expressed asn an "action"),
    # multiplied by a coefficient.
    hubbard_fermion = of.fermi_hubbard(
        x_dimension        = x_dimension,
        y_dimension        = y_dimension,
        tunneling          = tunneling,
        coulomb            = coulomb,
        chemical_potential = chemical_potential,
        magnetic_field     = magnetic_field,
        periodic           = periodic,
        spinless           = spinless)

    # Hubbard Hamiltonian expressed in QubitOperators, i.e. Pauli's (X, Y, Z) operators.
    # Each QubitOperator consists of the qubit it operates on (expressed as an "index")
    # and which Pauli is applied (expressed as an "action") multiplied by a coefficient.
    if _fermion_transform == 'bravyi-kitaev':
        n_qubits = x_dimension * y_dimension * (spinless and 1 or 2)
        hubbard_qubit = of.transforms.bravyi_kitaev(hubbard_fermion, n_qubits)
    elif _fermion_transform == 'jordan-wigner':
        hubbard_qubit = of.transforms.jordan_wigner(hubbard_fermion)

    # Remove terms below floating point epsilon
    hubbard_qubit.compress()

    return hubbard_qubit

def hamiltonian_matrix(x_dimension, y_dimension, tunneling, coulomb,
        chemical_potential = 0.00, magnetic_field = 0.0, periodic = True, spinless = False):
    """Create Fermi-Hubbard model Hamiltonian represented in matrix form"""

    hubbard_qubit = _hubbard_qubit(
        x_dimension        = x_dimension,
        y_dimension        = y_dimension,
        tunneling          = tunneling,
        coulomb            = coulomb,
        chemical_potential = chemical_potential,
        magnetic_field     = magnetic_field,
        periodic           = periodic,
        spinless           = spinless)

    return of.linalg.get_sparse_operator(hubbard_qubit).todense()

def hamiltonian_qiskit(x_dimension, y_dimension, tunneling, coulomb,
        chemical_potential = 0.00, magnetic_field = 0.0, periodic = True, spinless = False,
        fermion_transform='bravyi-kitaev'):
    """Create Fermi-Hubbard model Hamiltonian represented in matrix form"""

    hubbard_qubit = _hubbard_qubit(
        x_dimension        = x_dimension,
        y_dimension        = y_dimension,
        tunneling          = tunneling,
        coulomb            = coulomb,
        chemical_potential = chemical_potential,
        magnetic_field     = magnetic_field,
        periodic           = periodic,
        spinless           = spinless,
        fermion_transform  = fermion_transform)

    n_qubits = x_dimension * y_dimension * (spinless and 1 or 2)

    hubbard_qiskit = _to_qiskit(hubbard_qubit, n_qubits)

  # store the used fermion transform with the hamiltonian to ensure that
  # the objective function later uses the same transform
    hubbard_qiskit._fermion_transform = fermion_transform

    return hubbard_qiskit

class Model(object):
    """Convenience class to capture module parameters"""

    def __init__(self, xdim, ydim, t, U, chem=0.0, mag=0.0, periodic=True, spinless=False, precalc={}):
        self.x_dimension = xdim
        self.y_dimension = ydim
        self.tunneling   = t
        self.coulomb     = U
        self.chemical_potential = chem
        self.magnetic_field     = mag
        self.periodic = periodic
        self.spinless = spinless
        self._precalc = precalc

    def __call__(self):
        """Generate the model"""

        return self.x_dimension, self.y_dimension, self.tunneling, self.coulomb, \
               self.chemical_potential, self.magnetic_field, \
               self.periodic, self.spinless

    def initial(self, n_electrons_up, n_electrons_down, npar, transform='bravyi-kitaev', good=False):
        """\
        Provide a (good) initial and tight bounds for a given configuration

        Args:
            n_electrons_up(int): number of electrons with spin-up
            n_electrons_down(int): number of electrons with spin-down
            transform(str): for which fermion transform the initial applies
            good(bool): whether to return an initial close to the optimal

        Returns:
            initial(tuple): array of (good) initial parameters and an array of bounds
        """

        if good:
            at_opt = self.optimal(n_electrons_up, n_electrons_down, transform)
            if at_opt is not None:
                close = np.round(at_opt, npar <= 4 and 1 or 2)
                bounds = np.zeros((len(close), 2))
                bounds[:,0] = np.subtract(close, 0.1)
                bounds[:,1] = np.add(     close, 0.1)
                return close, bounds

        if npar <= 0:
            raise RuntimeError("not an optimizable configuration (%d parameters)" % npar)

        rng = np.random.default_rng(42)     # for reproducibility while debugging
        initial_amplitudes = np.array(-0.05+0.1*rng.random(size=npar))
        bounds = np.array([(-1.0, 1.0)]*npar)

        return initial_amplitudes, bounds

    def optimal(self, n_electrons_up, n_electrons_down, transform='bravyi-kitaev'):
        """\
        Lookup the pre-calculated optimal paramters

        Args:
            n_electrons_up(int): number of electrons with spin-up
            n_electrons_down(int): number of electrons with spin-down
            transform(str): for which fermion transform the initial applies

        Returns:
            optimum(tuple): array of parameters for the global minimum or None
        """

        try:
            return self._precalc[(n_electrons_up, n_electrons_down)]
        except KeyError:
            pass

        warnings.warn("No pre-calculated initial for configuration (%d, %d)" %\
                      (n_electrons_up, n_electrons_down))

        return None

small_model  = Model(2, 1, t=1.0, U=2.0,
    precalc={
        (1, 0) : np.array([-0.78536064,  0.89994575]),
        (0, 1) : np.array([-0.78536609, -0.25647772]),
        (1, 1) : np.array([-0.86866234,  0.18526051]),
    })

medium_model = Model(2, 2, t=1.0, U=2.0,
    precalc={
        (1, 1) : np.array([ 0.22048886, 0.22048479,  0.27563475,
                            0.22178354, 0.22177972,  0.24547588,
                            0.6276739,  0.60108877,  0.60108406]),
        (2, 2) : np.array([-0.81965099,  0.4858986 , -0.4858995,  0.76993761,
                            0.10298091, -0.03832318, -0.03832113, 0.64542339,
                            0.00399792, -0.00399722,  0.11716964, 0.32792626,
                           -0.06136483,  0.06136485]),
        (3, 3) : np.array([ 0.64501004, -0.6305074 , -0.63050858,
                            0.08473441,  0.06774534,  0.06774663,
                           -0.0411103 , -0.0411079 , -0.01508739])
    })

def exact(hubbard_hamiltonian, n_electrons_up, n_electrons_down):
    """Return the exact solution for the given Hubbard Model Hamiltonian"""

    n_qubits = hubbard_hamiltonian.num_qubits

  # JW has one-to-one mapping and thus the fermion operator's matrix can be used
  # directly; not so for other transformations
    spin_up_op   = sum([of.FermionOperator(((i, 1), (i, 0))) for i in range(0, n_qubits, 2)])
    spin_down_op = sum([of.FermionOperator(((i, 1), (i, 0))) for i in range(1, n_qubits, 2)])

    if hubbard_hamiltonian._fermion_transform == 'bravyi-kitaev':
        spin_up_op   = of.transforms.bravyi_kitaev(spin_up_op, n_qubits)
        spin_down_op = of.transforms.bravyi_kitaev(spin_down_op, n_qubits)
    elif hubbard_hamiltonian._fermion_transform == 'jordan-wigner':
        spin_up_op   = of.transforms.jordan_wigner(spin_up_op)
        spin_down_op = of.transforms.jordan_wigner(spin_down_op)

    spin_up_op   = _to_qiskit(spin_up_op, n_qubits)
    spin_down_op = _to_qiskit (spin_down_op, n_qubits)

  # get the matrix representations; note that the up_matrix will be 2x smaller if
  # n_qubits is even; same for down_matrix if odd
    hm_matrix   = hubbard_hamiltonian.to_matrix()
    up_matrix   = spin_up_op.to_matrix()
    down_matrix = spin_down_op.to_matrix()

    eigenvalues, eigenvectors = spla.eigh(hm_matrix)
    for i in range(hm_matrix.shape[0]):
        v = eigenvectors[:,i]
        n_up = float(np.real(v.T.dot(up_matrix).dot(v)))
        if not round(n_up, 3).is_integer() or not round(n_up) == n_electrons_up:
            continue

        n_down = float(np.real(v.T.dot(down_matrix).dot(v)))
        if not round(n_down, 3).is_integer() or not round(n_down) == n_electrons_down:
            continue

        return float(np.real(v.T.dot(hm_matrix).dot(v)))

    raise RuntimeError('configuration %d up, %d down not found' % (n_electrons_up, n_electrons_down))


import os.path
import logging
import numpy    as np
from   typing   import Union

import PyMieSim
from PyMieSim.Tools.LPModes     import GetFarFieldLP
from PyMieSim.Tools.Mesh        import FibonacciMesh
from PyMieSim.Tools.BaseClasses import BaseDetector, MeshProperty
from PyMieSim.Tools.ErrorMsg    import *
from PyMieSim.Tools.utils       import NA2Angle, LoadLPMode, IO
from PyMieSim.bin.LMTScatterer  import BindedPhotodiode, BindedLPMode

class Photodiode(BaseDetector, MeshProperty):
    """
    .. note::
        Detector type class representing a photodiode, light coupling is
        thus independant of the phase of the latter.

    Parameters
    ----------
    NA : :class:`float`
        Numerical aperture of imaging system.
    Sampling : :class:`int`
        Number of sampling points for the mode (inside NA).
    GammaOffset : :class:`float`
        Angle offset of detector in the direction perpendicular to polarization.
    PhiOffset : :class:`float`
        Angle offset of detector in the direction parallel to polarization.
    Filter : :class:`float`
        Angle of polarization filter in front of detector. Default is "None"
    CouplingMode : :class:`str`
        Methode for computing mode coupling. Either Point or Mean.

    """

    def __init__(self,
                 NA           : Union[int, float],
                 Sampling     : int                      = 400,
                 CouplingMode : str                      = 'Point',
                 GammaOffset  : Union[int, float]        = 0,
                 PhiOffset    : Union[int, float]        = 0,
                 Filter       : Union[int, float, bool]  = None):

        self.CouplingMode = CouplingMode
        self._Filter      = Filter
        self._PhiOffset   = PhiOffset
        self._GammaOffset = GammaOffset
        self._NA          = NA
        self._Sampling    = Sampling
        self.ScalarField  = np.ones(Sampling)

        self.GetBinding()

    def GetBinding(self):
        self.Mesh = FibonacciMesh(MaxAngle    = NA2Angle(self._NA).Radian,
                                  Sampling    = self._Sampling,
                                  PhiOffset   = self._PhiOffset,
                                  GammaOffset = self._GammaOffset)

        self.Bind = BindedPhotodiode(NA       = self._NA,
                                     Phi      = np.deg2rad(self._PhiOffset),
                                     Gamma    = np.deg2rad(self._GammaOffset),
                                     Filter   = np.deg2rad(self._Filter) if self._Filter else float(4242),
                                     Sampling = self._Sampling)

    def GetScalarField(self, Sampling, Structured=False):
        if Structured:
            return np.ones([Sampling, Sampling])
        else:
            return np.ones(Sampling)

    def __str__(self):
        return self.Name

    def __repr__(self):

        return IO( f"""
        Photodiode detector
        Coupling Mode: Intensity
        Numerical aperture:  {self.NA:.4f}
        Sampling:            {self.Mesh.Sampling}
        Gamma offset:        {self.Mesh.GammaOffset}
        Phi offset:          {self.Mesh.PhiOffset}
        """ )

class IntegratingSphere(Photodiode):
    """
    .. note::
        Detector type class representing a photodiode, light coupling is
        thus independant of the phase of the latter.

    Parameters
    ----------
    NA : :class:`float`
        Numerical aperture of imaging system.
    Sampling : :class:`int`
        Number of sampling points for the mode (inside NA).
    GammaOffset : :class:`float`
        Angle offset of detector in the direction perpendicular to polarization.
    PhiOffset : :class:`float`
        Angle offset of detector in the direction parallel to polarization.
    Filter : :class:`float`
        Angle of polarization filter in front of detector. Default is "None"
    CouplingMode : :class:`str`
        Methode for computing mode coupling. Either Point or Mean.

    """

    def __init__(self,
                 Sampling     : int                      = 400,
                 CouplingMode : str                      = 'Point',
                 Filter       : Union[int, float, bool]  = None):

        self.CouplingMode = CouplingMode
        self._Filter      = Filter
        self._PhiOffset   = 0
        self._GammaOffset = 0
        self._NA          = 2.0
        self._Sampling    = Sampling
        self.ScalarField  = np.ones(Sampling)

        self.GetBinding()

    def GetBinding(self):
        self.Mesh = FibonacciMesh(MaxAngle    = NA2Angle(self._NA).Radian,
                                  Sampling    = self._Sampling,
                                  PhiOffset   = self._PhiOffset,
                                  GammaOffset = self._GammaOffset)

        self.Bind = BindedPhotodiode(Sampling = self._Sampling,
                                     NA       = self._NA,
                                     Phi      = np.deg2rad(self._PhiOffset),
                                     Gamma    = np.deg2rad(self._GammaOffset),
                                     Filter   = np.deg2rad(self._Filter) if self._Filter else float(4242)
                                     )

    def GetScalarField(self, Sampling, Structured=False):
        if Structured:
            return np.ones([Sampling, Sampling])
        else:
            return np.ones(Sampling)

    def __repr__(self):

        return IO( f"""
        Integrating sphere
        Coupling Mode: Intensity
        Sampling:      {self.Mesh.Sampling}
        """ )

    def __str__(self):
        return self.Name

class LPmode(BaseDetector, MeshProperty):
    """
    .. note::
        Detector type class representing a fiber LP mode, light coupling is
        thus dependant of the phase of the latter.

    Parameters
    ----------
    Mode : :class:`tuple`
        LP mode index l, m.
    NA : :class:`float`
        Numerical aperture of imaging system.
    Sampling : :class:`int`
        Number of sampling points for the mode (inside NA).
    InterpSampling : :class:`int`
        Number of sampling point for interpolation of FarField mode.
    GammaOffset : :class:`float`
        Angle offset of detector in the direction perpendicular to polarization.
    PhiOffset : :class:`float`
        Angle offset of detector in the direction parallel to polarization.
    Filter : :class:`float`
        Angle of polarization filter in front of detector. Default is "None"
    CouplingMode : :class:`str`
        Methode for computing mode coupling. Either [Point or Mean].

    """

    def __init__(self,
                 Mode         : Union[tuple, list],
                 NA           : float,
                 Rotation     : Union[int, float]        = 0,
                 Sampling     : int                      = 401,
                 GammaOffset  : Union[int, float]        = 0,
                 PhiOffset    : Union[int, float]        = 0,
                 Filter       : Union[int, float, bool]  =  None,
                 CouplingMode : str                      = 'Point'):

        assert CouplingMode in ['Point','Mean'], Error_MeanCentered

        if NA > 1 or NA < 0: logging.warning(warning_NAHigh)

        self.ModeNumber   = Mode
        self.CouplingMode = CouplingMode
        self._Filter      = Filter
        self._PhiOffset   = PhiOffset
        self._GammaOffset = GammaOffset
        self._NA          = NA
        self._Sampling    = Sampling

        self.GetBinding()

    def GetBinding(self):
        self.Mesh = FibonacciMesh(MaxAngle    = NA2Angle(self._NA).Radian,
                                  Sampling    = self._Sampling,
                                  PhiOffset   = self._PhiOffset,
                                  GammaOffset = self._GammaOffset)

        self.ScalarField  = GetFarFieldLP(Mode     = self.ModeNumber,
                                          MaxAngle = NA2Angle(self._NA).Radian,
                                          Sampling = self._Sampling)

        self.Bind = BindedLPMode(ScalarField = self.ScalarField,
                                 Sampling    = self._Sampling,
                                 NA          = self._NA,
                                 Phi         = np.deg2rad(self._PhiOffset),
                                 Gamma       = np.deg2rad(self._GammaOffset),
                                 Filter      = np.deg2rad(self._Filter) if self._Filter else float(4242)
                                 )

    def GetScalarField(self, Sampling, Structured=False):

        return GetFarFieldLP(Mode       = self.ModeNumber,
                             MaxAngle   = NA2Angle(self._NA).Radian,
                             Sampling   = Sampling,
                             Structured = Structured)

    def __str__(self):
        return self.Name

    def __repr__(self):
        return IO( f"""
        LP mode detector
        Coupling Mode: Amplitude
        LP Mode:             {self.ModeNumber}
        Numerical aperture:  {self.NA:.4f}
        Sampling:            {self.Mesh.Sampling}
        Gamma offset:        {self.Mesh.GammaOffset}
        Phi offset:          {self.Mesh.PhiOffset}
        """ )

# -


"""Keep track of spacecraft jitter."""
import settings
import numpy as np
import astropy.table
import os.path
import zachopy.utils
import matplotlib.pylab as plt
import scipy.interpolate
import logging
import matplotlib.gridspec as gridspec
from settings import log_file_handler

logger = logging.getLogger(__name__)
logger.addHandler(log_file_handler)

def makeCartoon(seed=1):
    """Create a cartoon jitter timeseries"""

    np.random.seed(seed)
    rmsat2s = 2.0 / 3.0
    rmsat120s = 0.21

    # share across two dimensions
    rmsat120s1d = rmsat120s / np.sqrt(2)
    rmsat2s1d = rmsat2s / np.sqrt(2)

    cadence = 2.0
    smoothscale = 10.0
    nsmooth = int(smoothscale / cadence)
    t = np.arange(0, 30 * 24 * 60 * 60, 2)
    n = len(t)
    d = {}
    d['t'] = t
    for k in ['x', 'y']:
        v = np.random.normal(0, 1, n)
        for i in range(2):
            v = np.convolve(v, np.ones(nsmooth), mode='same')
        d[k] = v / np.std(v) * rmsat2s1d

    table = astropy.table.Table(d, names=['t', 'x', 'y'])
    table.write(os.path.join(settings.inputs, 'cartoon.jitter'), format='ascii.fixed_width', bookend=False, overwrite=True)

class Jitter(object):
    def __init__(self, camera=None, jitterrms=None, rawjitterbasename="AttErrTimeArcsec_80k.dat",
                 nsubpixelsperpixel=None,
                 amplifyinterexposurejitter=1.0):

        # set an extra directory specifically for this object
        self.directory = 'jitter/'

        # store the input camera
        self.camera = camera

        # set up the initial raw jitter file
        # (this one cam from Roland, some time ago)
        self.rawfile = os.path.join(settings.inputs, rawjitterbasename)

        # what do you want the RMS to be rescaled to?
        self.jitterrms = jitterrms

        # how much should the exposure to exposure jitter be amplified
        self.amplifyinterexposurejitter = amplifyinterexposurejitter

        # (for creating map that will be used to convolve with the psf)
        self.nsubpixelsperpixel = nsubpixelsperpixel

        # update the jitterball to one that has been binned to this cadence
        self.load()

    def load(self, remake=False):
        """make sure that a jitterball (timeseries of roll,pitch,yaw) has
      been loaded and binned to the appropriate exposure times"""

        try:
            # if the jitterball is already loaded into memory
            #  *and* of the correct cadence, we're all set!
            self.jitterball

            # make sure the we're using the right jitterball for this cadence
            assert (self.jittercadence == self.camera.cadence)

            # make sure we're not trying to remake the jitterball
            assert (remake == False)

        except (AttributeError, AssertionError):
            # load the processed jitterball
            self.loadProcessedJitterball()

    @property
    def basename(self):
        cadencestatement = '.cadence{:.0f}s'.format(self.camera.cadence)

        if self.jitterrms is not None:
            jitterstatement = '.rescaledto{:0.2f}arcsec'.format(self.jitterrms)
        else:
            jitterstatement = '.unscaled'

        return os.path.basename(self.rawfile) + cadencestatement + jitterstatement

    @property
    def processedfile(self):
        """determine what the processed filename should be (based on rawfile)"""

        # store the processed files in the intermediates directory
        directory = settings.intermediates + self.directory

        # make sure a jitter directory actually exists
        zachopy.utils.mkdir(directory)

        # define the filename
        return os.path.join(directory, self.basename + '.processed.npy')

    def loadProcessedJitterball(self):
        """load a pre-processed jitterball, from the intermediates directory"""

        # if not, populate the jitterball for this cadence
        logger.info(
            'populating the jitterball for {0:.0f} second cadence, '
            'based on the raw jitter file {1}.'.format(
                self.camera.cadence,
                self.basename))

        try:
            # if a processed file already exists, load it
            self.jitterball, self.jittermap = np.load(self.processedfile)
            self.jittercadence = self.camera.cadence
        except IOError:

            logger.info('no processed jitter file was found for {}'.format(
                self.basename))

            self.loadUnprocessedJitterball()

    def loadUnprocessedJitterball(self):
        """load from a raw jitter file, process, and save"""

        # otherwise, create a binned jitter structure
        logger.info('loading raw jitter from {}'.format(self.rawfile))

        # load the raw file
        if 'AttErrTimeArcsec' in self.rawfile:
            self.rawdata = astropy.io.ascii.read(self.rawfile,
                                                 names=['t', 'x', 'y', 'z'])
        else:
            self.rawdata = astropy.io.ascii.read(self.rawfile, names=['t', 'x', 'y'])

        # subtract means
        self.rawdata['x'] -= np.mean(self.rawdata['x'])
        self.rawdata['y'] -= np.mean(self.rawdata['y'])

        # scale jitterball to requirements (should be inflation by ~1.5)
        if self.jitterrms is not None:
            # STILL A KLUDGE! NEED ROLL, PITCH, YAW!
            original_rms = np.sqrt(np.mean(self.rawdata['x'] ** 2 +
                                           self.rawdata['y'] ** 2))
            self.rawdata['x'] *= self.jitterrms / original_rms
            self.rawdata['y'] *= self.jitterrms / original_rms

        # smooth them to the required cadence
        logger.info("smoothing the jitter to {0}s cadence".format(
            self.camera.cadence))

        #  figure out the time-spacing of the jitter timeseries
        spacings = self.rawdata['t'][1:] - self.rawdata['t'][:-1]
        spacing = np.median(spacings)

        # make sure that the jitter timeseries is evenly spaced
        aboutright = 0.01
        assert ((np.abs(spacings - spacing) < spacing * aboutright).all())

        # create a convolution filter, to smooth to camera's cadence
        n = np.long(self.camera.cadence / spacing)
        filter = np.ones(n) / n

        # construct smoothed timeseries, sampled at raw time resolution
        smoothed_t = np.convolve(self.rawdata['t'], filter, mode='valid')
        smoothed_x = np.convolve(self.rawdata['x'], filter, mode='valid')
        smoothed_y = np.convolve(self.rawdata['y'], filter, mode='valid')

        # sample smoothed timeseries at the camera's cadence
        t = smoothed_t[::n]
        x = smoothed_x[::n]
        y = smoothed_y[::n]

        # plot each dimension separately
        logger.info('saving binned jitter timeseries plot')

        # create the plot of the timeseries
        plotdirectory = os.path.join(settings.plots, self.directory)
        zachopy.utils.mkdir(plotdirectory)
        bkw = dict(alpha=0.5, color='black')
        rkw = dict(linewidth=2, alpha=0.5, marker='o', color='red')
        fi, ax = plt.subplots(2, 1, sharey=True, sharex=True)
        ax[0].plot(self.rawdata['t'], self.rawdata['x'], **bkw)
        ax[0].plot(t, x, **rkw)
        ax[1].plot(self.rawdata['t'], self.rawdata['y'], **bkw)
        ax[1].plot(t, y, **rkw)
        ax[0].set_xlim(0, self.camera.cadence * 10)
        ax[0].set_title('TESS Pointing Jitter for \n{}\nfor {}s Cadence'.format(self.basename, self.camera.cadence),
                        fontsize=6)
        ax[0].set_ylabel('x (")')
        ax[1].set_ylabel('y (")')
        ax[1].set_xlabel('Time (seconds)')
        fi.savefig(os.path.join(plotdirectory,
                                self.basename + '_timeseries.pdf'))

        # make interpolators to keep track of the running smooth means
        ikw = dict(kind='nearest', fill_value=0, bounds_error=False)
        xip = scipy.interpolate.interp1d(smoothed_t, smoothed_x, **ikw)
        yip = scipy.interpolate.interp1d(smoothed_t, smoothed_y, **ikw)

        # assign the jittermap here, in units of subpixels, to be used for convolution in the PSF code
        arcsectosubpixels = 1.0 / self.camera.pixelscale * self.nsubpixelsperpixel
        xoff = (self.rawdata['x'] - xip(self.rawdata['t'])) * arcsectosubpixels
        yoff = (self.rawdata['y'] - yip(self.rawdata['t'])) * arcsectosubpixels

        npixelsfromcenter = 1
        nbins = np.maximum(np.max(np.abs(np.sqrt(xoff ** 2 + yoff ** 2))),
                           1)  # npixelsfromcenter*self.nsubpixelsperpixel
        limits = [[-nbins, nbins], [-nbins, nbins]]

        # define the jittermap as a 2D histrogram for convolution within exps
        self.jittermap = np.histogram2d(xoff, yoff,
                                        bins=nbins,
                                        range=limits,
                                        normed=True)

        # define the binned jitterball, for nudges between exps
        self.jitterball = (x, y)

        # keep track of the jitter cadence associated with this
        self.jittercadence = self.camera.cadence

        logger.info('saving jittermap plots')

        # plot the adopted jitterball, as more useful binning
        plothist2d(self.jittermap, scale=1.0 / self.nsubpixelsperpixel,
                   title='TESS Pointing Jitter over {0}s'.format(
                       self.camera.cadence),
                   xtitle='Pixels', ytitle='Pixels',
                   filename=os.path.join(plotdirectory, self.basename + '_jittermap.pdf'))

        # save the necessary jitter files
        logger.info('saving the jitter files to {0}'.format(self.processedfile))
        np.save(self.processedfile, (self.jitterball, self.jittermap))

    @property
    def x(self):
        return self.amplifyinterexposurejitter * self.jitterball[0]

    @property
    def y(self):
        return self.amplifyinterexposurejitter * self.jitterball[1]

    def writeNudges(self, outfile='jitter.txt'):

        counters = np.arange(len(self.x))
        bjds = self.camera.counterToBJD(counters)
        time = bjds - np.min(bjds)
        plt.figure('jitter timeseries')
        gs = gridspec.GridSpec(2, 1, hspace=0.15)
        kw = dict(linewidth=2)
        ax = None

        for i, what in enumerate((self.x, self.y)):
            ax = plt.subplot(gs[i], sharex=ax, sharey=ax)
            ax.plot(time, what, **kw)
            ax.set_ylabel(['dRA (arcsec)', 'dDec (arcsec)'][i])
            if i == 0:
                ax.set_title('Jitter Timeseries from\n{}'.format(self.basename))

        plt.xlabel('Time from Observation Start (days)')
        plt.xlim(np.min(time), np.max(time))
        plt.draw()
        plt.savefig(outfile.replace('.txt', '.pdf'))

        data = [counters, bjds, self.x, self.y]
        names = ['imagenumber', 'bjd', 'arcsecnudge_ra', 'arcsecnudge_dec']

        t = astropy.table.Table(data=data, names=names)
        t.write(outfile.replace('.txt', '_amplifiedby{}.txt'.format(self.amplifyinterexposurejitter)),
                format='ascii.fixed_width', delimiter=' ', overwrite=True)
        logger.info("save jitter nudge timeseries to {0}".format(outfile))

    def applyNudge(self,
                   counter=None,  # which row to use from jitterball?
                   dx=None, dy=None,  # custom nudges, in arcsec
                   header=None,  # the FITS header in which to record nudges
                   ):

        """jitter the camera by a little bit,
      by introducing nudges draw from a
      (cadence-appropriate) jitterball timeseries."""

        # make sure the jitterball has been populated
        self.load()
        n = len(self.x)

        # should we be applying a custom offset?
        usecustom = (counter is None)
        if usecustom:
            self.camera.nudge['x'] = dx
            self.camera.nudge['y'] = dy
        else:
            # if we're over the counter, loop back
            i = counter % n
            self.camera.nudge['x'] = self.x[i]
            self.camera.nudge['y'] = self.y[i]

        # if possible, write the details to the supplied FITS header
        try:
            header['MOTION'] = ''
            header['MOTNOTE'] = ('',
                                 'properties of the image motion applied')
            header['JITTERX'] = (self.camera.nudge['x'],
                                 '["] jitter-induced nudge')
            header['JITTERY'] = (self.camera.nudge['y'],
                                 '["] jitter-induced nudge')
            header['JITPFILE'] = (self.basename,
                                  'processed jitter filename')
            header['JITSCALE'] = (self.amplifyinterexposurejitter, 'jitter magnified by ? relative to file')
            header['JITCOUNT'] = (i, 'which row of jitter file was applied?')

            logger.info('updated header keywords')
        except TypeError:
            logger.info('no header was found to update')

        # move the camera, using the updated nudge values
        logger.info("nudged the camera to {x},{y}"
                   " away from nominal pointing.".format(**self.camera.nudge))

def plothist2d(hist, title=None, log=False, scale=1.0,
               xtitle=None, ytitle=None, filename=None):
    """Plot a 2D histogram."""
    map = hist[0]
    x = (hist[1][1:] + (hist[1][0] - hist[1][1]) / 2.0) * scale
    y = (hist[2][1:] + (hist[2][0] - hist[2][1]) / 2.0) * scale
    fig = plt.figure(figsize=(10, 10))
    plt.clf()
    plt.subplots_adjust(hspace=0, wspace=0)
    ax_map = fig.add_subplot(2, 2, 3)
    ax_vert = fig.add_subplot(2, 2, 4, sharey=ax_map)
    ax_hori = fig.add_subplot(2, 2, 1, sharex=ax_map)

    ax_hori.plot(x, np.sum(map, 0) / np.sum(map), marker='o', color='black', linewidth=3)
    ax_vert.plot(np.sum(map, 1) / np.sum(map), y, marker='o', color='black', linewidth=3)
    if log:
        ax_vert.semilogx()
        ax_hori.semilogy()
    if log:
        bottom = np.min(map[map > 0]) / np.maximum(np.sum(map, 0).max(), np.sum(map, 1).max())
    else:
        bottom = 0
    top = 1
    ax_hori.set_ylim(bottom, top)
    ax_vert.set_xlim(bottom, top)

    ax_vert.tick_params(labelleft=False)
    ax_hori.tick_params(labelbottom=False)
    if title is not None:
        ax_hori.set_title(title)
    if xtitle is not None:
        ax_map.set_xlabel(xtitle)
    if ytitle is not None:
        ax_map.set_ylabel(ytitle)

    try:
        xhalf, yhalf = (x[1] - x[0]) / 2.0, (y[1] - y[0]) / 2.0
    except IndexError:
        xhalf, yhalf = 0.5, 0.5

    kw = dict(cmap='gray_r',
              extent=[x.min() - xhalf, x.max() + xhalf,
                      y.min() - yhalf, y.max() + yhalf],
              interpolation='nearest')
    if log:
        y = np.log(map)
    else:
        y = map
    ax_map.imshow(y, **kw)
    if filename is not None:
        fig.savefig(filename)

import numpy,import scipy,import astropy


"""Implement the GeometricProgram class"""
from __future__ import print_function
# unicode_literals here interfere with the boundschecking example
import sys
from time import time
from collections import defaultdict
import numpy as np
from ..nomials import NomialData
from ..small_classes import CootMatrix, SolverLog, Numbers, FixedScalar
from ..keydict import KeyDict
from ..small_scripts import mag
from ..solution_array import SolutionArray
from .costed import CostedConstraintSet
from ..exceptions import InvalidPosynomial

DEFAULT_SOLVER_KWARGS = {"cvxopt": {"kktsolver": "ldl"}}
SOLUTION_TOL = {"cvxopt": 1e-3, "mosek_cli": 1e-4, "mosek": 1e-5,
                'mosek_conif': 1e-3}

def _get_solver(solver, kwargs):
    """Get the solverfn and solvername associated with solver"""
    if solver is None:
        from .. import settings
        solver = settings.get("default_solver", None)
        if not solver:
            raise ValueError(
                "No solver was given; perhaps gpkit was not properly"
                " installed, or found no solvers during the"
                " installation process.")

    if solver == "cvxopt":
        from .._cvxopt import cvxoptimize
        solverfn = cvxoptimize
    elif solver == "mosek_cli":
        from .._mosek import cli_expopt
        solverfn = cli_expopt.imize_fn(**kwargs)
    elif solver == "mosek":
        from .._mosek import expopt
        solverfn = expopt.imize
    elif solver == 'mosek_conif':
        from .._mosek import mosek_conif
        solverfn = mosek_conif.mskoptimize
    elif hasattr(solver, "__call__"):
        solverfn = solver
        solver = solver.__name__
    else:
        raise ValueError("Unknown solver '%s'." % solver)
    return solverfn, solver

class GeometricProgram(CostedConstraintSet, NomialData):
    # pylint: disable=too-many-instance-attributes
    """Standard mathematical representation of a GP.

    Attributes with side effects
    ----------------------------
    `solver_out` and `solver_log` are set during a solve
    `result` is set at the end of a solve if solution status is optimal

    Examples
    --------
    >>> gp = gpkit.geometric_program.GeometricProgram(
                        # minimize
                        x,
                        [   # subject to
                            1/x  # <= 1, implicitly
                        ], {})
    >>> gp.solve()
    """
    def __init__(self, cost, constraints, substitutions,
                 allow_missingbounds=False):
        # pylint:disable=super-init-not-called
        # initialize attributes modified by internal methods
        self._result = None
        self.v_ss = None
        self.nu_by_posy = None
        self.solver_log = None
        self.solver_out = None
        self.__bare_init__(cost, constraints, substitutions, varkeys=False)
        for key, sub in self.substitutions.items():
            if isinstance(sub, FixedScalar):
                sub = sub.value
                if hasattr(sub, "units"):
                    sub = sub.to(key.units or "dimensionless").magnitude
                self.substitutions[key] = sub
            # only allow Numbers and ndarrays
            if not isinstance(sub, (Numbers, np.ndarray)):
                raise ValueError("substitution {%s: %s} with value type %s is"
                                 " not allowed in .substitutions."
                                 % (key, sub, type(sub)))
        try:
            self.posynomials = [cost.sub(self.substitutions)]
        except InvalidPosynomial:
            raise InvalidPosynomial("cost must be a Posynomial")
        self.posynomials.extend(self.as_posyslt1(self.substitutions))
        self.hmaps = [p.hmap for p in self.posynomials]
        ## Generate various maps into the posy- and monomials
        # k [j]: number of monomials (rows of A) present in each constraint
        self.k = [len(hm) for hm in self.hmaps]
        p_idxs = []  # p_idxs [i]: posynomial index of each monomial
        self.m_idxs = []  # m_idxs [i]: monomial indices of each posynomial
        for i, p_len in enumerate(self.k):
            self.m_idxs.append(list(range(len(p_idxs), len(p_idxs) + p_len)))
            p_idxs += [i]*p_len
        self.p_idxs = np.array(p_idxs)
        # m_idxs: first exp-index of each monomial equality
        self.meq_idxs = {sum(self.k[:i]) for i, p in enumerate(self.posynomials)
                         if getattr(p, "from_meq", False)}
        self.gen()  # A [i, v]: sparse matrix of powers in each monomial
        if self.missingbounds and not allow_missingbounds:
            boundstrs = "\n".join("  %s has no %s bound%s" % (v, b, x)
                                  for (v, b), x in self.missingbounds.items())
            raise ValueError("Geometric Program is not fully bounded:\n"
                             + boundstrs)

    varkeys = NomialData.varkeys

    def gen(self):
        "Generates nomial and solve data (A, p_idxs) from posynomials"
        self._hashvalue = self._varlocs = self._varkeys = None
        self._exps, self._cs = [], []
        for hmap in self.hmaps:
            self._exps.extend(hmap.keys())
            self._cs.extend(hmap.values())
        self.vks = self.varlocs
        self.A, self.missingbounds = genA(self.exps, self.varlocs,
                                          self.meq_idxs)

    # pylint: disable=too-many-statements, too-many-locals
    def solve(self, solver=None, verbosity=1, warn_on_check=False,
              process_result=True, gen_result=True, **kwargs):
        """Solves a GeometricProgram and returns the solution.

        Arguments
        ---------
        solver : str or function (optional)
            By default uses one of the solvers found during installation.
            If set to "mosek", "mosek_cli", or "cvxopt", uses that solver.
            If set to a function, passes that function cs, A, p_idxs, and k.
        verbosity : int (default 1)
            If greater than 0, prints solver name and solve time.
        **kwargs :
            Passed to solver constructor and solver function.

        Returns
        -------
        result : SolutionArray
        """
        solverfn, solvername = _get_solver(solver, kwargs)

        starttime = time()
        if verbosity > 0:
            print("Using solver '%s'" % solvername)
            print("Solving for %i variables." % len(self.varlocs))

        solver_kwargs = DEFAULT_SOLVER_KWARGS.get(solvername, {})
        solver_kwargs.update(kwargs)

        # NOTE: SIDE EFFECTS AS WE LOG SOLVER'S STDOUT AND OUTPUT
        original_stdout = sys.stdout
        self.solver_log = SolverLog(verbosity-1, original_stdout)
        try:
            sys.stdout = self.solver_log   # CAPTURED
            solver_out = solverfn(c=self.cs, A=self.A, p_idxs=self.p_idxs,
                                  k=self.k, **solver_kwargs)
            self.solver_out = solver_out
        finally:
            sys.stdout = original_stdout
        # STDOUT HAS BEEN RETURNED. ENDING SIDE EFFECTS.
        self.solver_log = "\n".join(self.solver_log)

        solver_out["solver"] = solvername
        solver_out["soltime"] = time() - starttime
        if verbosity > 0:
            print("Solving took %.3g seconds." % (solver_out["soltime"],))

        solver_status = str(solver_out.get("status", None))
        if solver_status.lower() != "optimal":
            raise RuntimeWarning(
                "final status of solver '%s' was '%s', not 'optimal'.\n\n"
                "The solver's result is stored in model.program.solver_out. "
                "A result dict can be generated via "
                "program._compile_result(program.solver_out)." %
                (solvername, solver_status))

        if gen_result:  # NOTE: SIDE EFFECTS
            self._result = self.generate_result(solver_out, warn_on_check,
                                                verbosity, process_result)
            return self.result

        solver_out["gen_result"] = \
            lambda: self.generate_result(solver_out, dual_check=False)
        return solver_out

    @property
    def result(self):
        "Creates and caches a result from the raw solver_out"
        if not self._result:
            self._result = self.generate_result(self.solver_out)
        return self._result

    def generate_result(self, solver_out, warn_on_check=True, verbosity=0,
                        process_result=True, dual_check=True):
        "Generates a full SolutionArray and checks it."
        if verbosity > 1:
            tic = time()
        soltime = solver_out["soltime"]
        result = self._compile_result(solver_out)  # NOTE: SIDE EFFECTS
        if verbosity > 1:
            print("result packing took %.2g%% of solve time" %
                  ((time() - tic) / soltime * 100))
            tic = time()

        try:
            tol = SOLUTION_TOL.get(solver_out["solver"], 1e-5)
            self.check_solution(result["cost"], solver_out['primal'],
                                solver_out["nu"], solver_out["la"], tol)
        except RuntimeWarning as e:
            if warn_on_check:
                e = str(e)
                if dual_check or ("Dual" not in e and "nu" not in e):
                    print("Solution check warning: %s" % e)
            else:
                raise e
        if verbosity > 1:
            print("solution checking took %.2g%% of solve time" %
                  ((time() - tic) / soltime * 100))
            tic = time()

        if process_result:
            self.process_result(result)
        if verbosity > 1:
            print("processing results took %.2g%% of solve time" %
                  ((time() - tic) / soltime * 100))
        return result

    def _generate_nula(self, solver_out):
        if "nu" in solver_out:
            # solver gave us monomial sensitivities, generate posynomial ones
            nu = np.ravel(solver_out["nu"])
            self.nu_by_posy = [nu[mi] for mi in self.m_idxs]
            la = np.array([sum(nup) for nup in self.nu_by_posy])
        elif "la" in solver_out:
            # solver gave us posynomial sensitivities, generate monomial ones
            la = np.ravel(solver_out["la"])
            if len(la) == len(self.hmaps) - 1:
                # assume the solver dropped the cost's sensitivity (always 1.0)
                la = np.hstack(([1.0], la))
            Ax = np.ravel(self.A.dot(solver_out['primal']))
            z = Ax + np.log(self.cs)
            m_iss = [self.p_idxs == i for i in range(len(la))]
            self.nu_by_posy = [la[p_i]*np.exp(z[m_is])/sum(np.exp(z[m_is]))
                               for p_i, m_is in enumerate(m_iss)]
            nu = np.hstack(self.nu_by_posy)
        else:
            raise RuntimeWarning("The dual solution was not returned.")
        solver_out["nu"], solver_out["la"] = nu, la

    def _compile_result(self, solver_out):
        """Creates a result dict (as returned by solve() from solver output

        This internal method is called from within the solve() method, unless
        solver_out["status"] is not "optimal", in which case a RuntimeWarning
        is raised prior to this method being called. In that case, users
        may use this method to attempt to create a results dict from the
        output of the failed solve.

        Arguments
        ---------
        solver_out: dict
            dict in format returned by solverfn within GeometricProgram.solve

        Returns
        -------
        result: dict
            dict in format returned by GeometricProgram.solve()
        """
        self._generate_nula(solver_out)
        primal = solver_out["primal"]
        nu, la = solver_out["nu"], solver_out["la"]
        # confirm lengths before calling zip
        if not self.varlocs and len(primal) == 1 and primal[0] == 0:
            primal = []  # an empty result, as returned by MOSEK
        assert len(self.varlocs) == len(primal)
        result = {"freevariables": KeyDict(zip(self.varlocs, np.exp(primal)))}
        # get cost #
        if "objective" in solver_out:
            result["cost"] = float(solver_out["objective"])
        else:
            # use self.posynomials[0] because the cost may have had constants
            freev = result["freevariables"]
            cost = self.posynomials[0].sub(freev)
            if cost.varkeys:
                raise ValueError("cost contains unsolved variables %s"
                                 % cost.varkeys.keys())
            result["cost"] = mag(cost.c)
        # get variables #
        result["constants"] = KeyDict(self.substitutions)
        result["variables"] = KeyDict(result["freevariables"])
        result["variables"].update(result["constants"])
        # get sensitivities #
        result["sensitivities"] = {"nu": nu, "la": la}
        self.v_ss = self.sens_from_dual(la[1:].tolist(), self.nu_by_posy[1:],
                                        result)
        # add cost's sensitivity in (nu could be self.nu_by_posy[0])
        cost_senss = {var: sum([self.cost.exps[i][var]*nu[i] for i in locs])
                      for (var, locs) in self.cost.varlocs.items()}
        var_senss = self.v_ss.copy()
        for key, value in cost_senss.items():
            var_senss[key] = value + var_senss.get(key, 0)
        # carry linked sensitivities over to their constants
        for v in list(v for v in var_senss if v.gradients):
            dlogcost_dlogv = var_senss.pop(v)
            val = result["constants"][v]
            for c, dv_dc in v.gradients.items():
                if val != 0:
                    dlogv_dlogc = dv_dc * result["constants"][c]/val
                # make nans / infs explicitly to avoid warnings
                elif dlogcost_dlogv == 0:
                    dlogv_dlogc = np.nan
                else:
                    dlogv_dlogc = np.inf * dv_dc*result["constants"][c]
                accum = var_senss.get(c, 0)
                var_senss[c] = dlogcost_dlogv*dlogv_dlogc + accum
                if v in cost_senss:
                    if c in self.cost.varkeys:
                        dlogcost_dlogv = cost_senss.pop(v)
                        accum = cost_senss.get(c, 0)
                        cost_senss[c] = dlogcost_dlogv*dlogv_dlogc + accum

        result["sensitivities"]["cost"] = cost_senss
        result["sensitivities"]["variables"] = KeyDict(var_senss)
        result["sensitivities"]["constants"] = KeyDict(
            {k: v for k, v in var_senss.items() if k in result["constants"]})
        result["soltime"] = solver_out["soltime"]
        return SolutionArray(result)

    def check_solution(self, cost, primal, nu, la, tol, abstol=1e-20):
        """Run a series of checks to mathematically confirm sol solves this GP

        Arguments
        ---------
        cost:   float
            cost returned by solver
        primal: list
            primal solution returned by solver
        nu:     numpy.ndarray
            monomial lagrange multiplier
        la:     numpy.ndarray
            posynomial lagrange multiplier

        Raises
        ------
        RuntimeWarning, if any problems are found
        """
        def _almost_equal(num1, num2):
            "local almost equal test"
            return (num1 == num2 or abs((num1 - num2) / (num1 + num2)) < tol
                    or abs(num1 - num2) < abstol)
        A = self.A.tocsr()
        # check primal sol
        primal_exp_vals = self.cs * np.exp(A.dot(primal))   # c*e^Ax
        if not _almost_equal(primal_exp_vals[self.m_idxs[0]].sum(), cost):
            raise RuntimeWarning("Primal solution computed cost did not match"
                                 " solver-returned cost: %s vs %s" %
                                 (primal_exp_vals[self.m_idxs[0]].sum(), cost))
        for mi in self.m_idxs[1:]:
            if primal_exp_vals[mi].sum() > 1 + tol:
                raise RuntimeWarning("Primal solution violates constraint:"
                                     " %s is greater than 1." %
                                     primal_exp_vals[mi].sum())
        # check dual sol
        # note: follows dual formulation in section 3.1 of
        # http://web.mit.edu/~whoburg/www/papers/hoburg_phd_thesis.pdf
        if not _almost_equal(self.nu_by_posy[0].sum(), 1.):
            raise RuntimeWarning("Dual variables associated with objective sum"
                                 " to %s, not 1" % self.nu_by_posy[0].sum())
        if any(nu < 0):
            if all(nu > -tol/1000.):  # HACK, see issue 528
                print("Allowing negative dual variable(s) as small as"
                      " %s." % min(nu))
            else:
                raise RuntimeWarning("Dual solution has negative entries as"
                                     " small as %s." % min(nu))
        ATnu = A.T.dot(nu)
        if any(np.abs(ATnu) > tol):
            raise RuntimeWarning("sum of nu^T * A did not vanish")
        b = np.log(self.cs)
        dual_cost = sum(
            self.nu_by_posy[i].dot(
                b[mi] - np.log(self.nu_by_posy[i]/la[i]))
            for i, mi in enumerate(self.m_idxs) if la[i])
        if not _almost_equal(np.exp(dual_cost), cost):
            raise RuntimeWarning("Dual cost %s does not match primal"
                                 " cost %s" % (np.exp(dual_cost), cost))

def genA(exps, varlocs, meq_idxs):  # pylint: disable=invalid-name
    """Generates A matrix

    Returns
    -------
        A : sparse Cootmatrix
            Exponents of the various free variables for each monomial: rows
            of A are monomials, columns of A are variables.
        missingbounds : dict
            Keys: variables that lack bounds. Values: which bounds are missed.
    """
    missingbounds = {}
    row, col, data = [], [], []
    for j, var in enumerate(varlocs):
        upperbound, lowerbound = False, False
        row.extend(varlocs[var])
        col.extend([j]*len(varlocs[var]))
        data.extend(exps[i][var] for i in varlocs[var])
        for i in varlocs[var]:
            if i not in meq_idxs:
                if upperbound and lowerbound:
                    break
                elif exps[i][var] > 0:  # pylint:disable=simplifiable-if-statement
                    upperbound = True
                else:
                    lowerbound = True
        if not upperbound:
            missingbounds[(var, "upper")] = ""
        if not lowerbound:
            missingbounds[(var, "lower")] = ""

    check_mono_eq_bounds(missingbounds, gen_mono_eq_bounds(exps, meq_idxs))

    # space the matrix out for trailing constant terms
    for i, exp in enumerate(exps):
        if not exp:
            row.append(i)
            col.append(0)
            data.append(0)
    A = CootMatrix(row, col, data)

    return A, missingbounds

def gen_mono_eq_bounds(exps, meq_idxs):  # pylint: disable=too-many-locals
    "Generate conditional monomial equality bounds"
    meq_bounds = defaultdict(set)
    for i in meq_idxs:
        if i % 2:  # skip the second index of a meq
            continue
        p_upper, p_lower, n_upper, n_lower = set(), set(), set(), set()
        for key, x in exps[i].items():
            if x > 0:
                p_upper.add((key, "upper"))
                p_lower.add((key, "lower"))
            else:
                n_upper.add((key, "upper"))
                n_lower.add((key, "lower"))
        # (consider x*y/z == 1)
        # for a var (e.g. x) to be upper bounded by this monomial equality,
        #   - vars of the same sign/side (y) must be lower bounded
        #   - AND vars of the opposite sign/side (z) must be upper bounded
        p_ub = frozenset(n_upper).union(p_lower)
        p_lb = frozenset(n_lower).union(p_upper)
        n_ub = frozenset(p_upper).union(n_lower)
        n_lb = frozenset(p_lower).union(n_upper)
        for keys, ub, lb in ((p_upper, p_ub, p_lb), (n_upper, n_ub, n_lb)):
            for key, _ in keys:
                meq_bounds[(key, "upper")].add(ub.difference([(key, "lower")]))
                meq_bounds[(key, "lower")].add(lb.difference([(key, "upper")]))
    return meq_bounds

def check_mono_eq_bounds(missingbounds, meq_bounds):
    "Bounds variables with monomial equalities"
    still_alive = True
    while still_alive:
        still_alive = False  # if no changes are made, the loop exits
        for bound in list(meq_bounds):
            if bound not in missingbounds:
                del meq_bounds[bound]
                continue
            conditions = meq_bounds[bound]
            for condition in conditions:
                if not any(bound in missingbounds for bound in condition):
                    del meq_bounds[bound]
                    del missingbounds[bound]
                    still_alive = True
                    break
    for (var, bound) in meq_bounds:
        boundstr = (", but would gain it from any of these"
                    " sets of bounds: ")
        for condition in list(meq_bounds[(var, bound)]):
            meq_bounds[(var, bound)].remove(condition)
            newcond = condition.intersection(missingbounds)
            if newcond and not any(c.issubset(newcond)
                                   for c in meq_bounds[(var, bound)]):
                meq_bounds[(var, bound)].add(newcond)
        boundstr += " or ".join(str(list(condition))
                                for condition in meq_bounds[(var, bound)])
        missingbounds[(var, bound)] = boundstr