"""Sound Heat Map generation and distribution."""

import ast
import asyncio
import json
import logging
import optparse
import os
import re
import sys
import traceback
from datetime import datetime, timedelta

import coloredlogs
import dateutil.parser as dp
import numpy as np
import pandas as pd
import requests
from beeprint import pp
from hbmqtt.client import ClientException, MQTTClient
from hbmqtt.mqtt.constants import QOS_1
from pyost import SensorThingStore
from shapely.geometry import Point
from shapely.geometry.polygon import Polygon
from model import calculate_shm_dir, grid_area
from weighting import weightfunc

logging.getLogger("urllib3").setLevel(logging.WARNING)  # not from these
logging.getLogger("requests").setLevel(logging.WARNING)
logging.getLogger("asyncio").setLevel(logging.WARNING)
logging.getLogger("hbmqtt").setLevel(logging.WARNING)
logging.getLogger("transitions").setLevel(logging.WARNING)

FREQUENCIES = np.array(ast.literal_eval(os.environ["FREQUENCIES"]))
MQTT_BROKER_PREFIX = os.environ["MQTT_BROKER_PREFIX"]
GOST_URL = os.environ["GOST_URL"]
UPDATE_DATASTREAMS_PERIODE = float(os.environ["UPDATE_DATASTREAMS_PERIODE"])
NEW_MAP_PERIODE = float(os.environ["NEW_MAP_PERIODE"])
MAP_LATENCY = float(os.environ["MAP_LATENCY"])
MAP_AVERAGING_TIME = float(os.environ["MAP_AVERAGING_TIME"])
FILTER_MICS = os.environ["FILTER_MICS"]
WEIGHTINGS = ast.literal_eval(os.environ["WEIGHTINGS"])

try:
    SHM_ID = os.environ["SHM_ID"]
except KeyError:
    # No instance set
    SHM_ID = None

try:
    AUTH = (os.environ["GOST_USER"], os.environ["GOST_PASS"])
except KeyError:
    AUTH = None

def CPBLZeq_Observation_to_dataframe(message):
    """Convert MQTT CPBLZeq observation message to a time indexable Dataframe."""
    # multiple observations can be bundled in one message
    observations = json.loads(message.data)["result"]["response"]["value"]

    all_values = []
    all_times = []
    for obs in observations:
        values = np.array(obs["values"])
        nmeas, nfreq = values.shape
        assert nfreq == 33

        start = dp.parse(obs["startTime"])
        end = dp.parse(obs["endTime"])

        times = start + (end - start) / nmeas * np.arange(nmeas)

        all_values.append(values)
        all_times.append(times)

    all_values = np.concatenate(all_values, axis=0)
    all_times = np.concatenate(all_times)

    return pd.DataFrame(all_values, index=all_times, columns=FREQUENCIES)

class Datastream:
    """Object representation of a Datastream."""

    def __init__(self, json):
        self.id = json["@iot.id"]
        self.location = json["observedArea"]["coordinates"]
        # coordinates must be given in (x=lon, y=lat) format, but we use lat/lon intern.
        self.location = self.location[::-1]
        self.description = json["description"]
        self.name = json["name"]
        self.mqtt_topic = MQTT_BROKER_PREFIX + f"/Datastreams({self.id})/Observations"
        self.recent_data = None

    def add(self, data):
        self.recent_data = pd.concat((self.recent_data, data))

    def average_and_drop_old_data(self, starttime, endtime):
        # average
        if self.recent_data is None or self.recent_data.empty:
            return None

        avg = self.recent_data[starttime:endtime].mean(axis=0)

        logging.debug(
            "datastream.recent_data before drop:" + pp(self.recent_data, output=False)
        )

        # drop data that will never be used again
        self.recent_data = self.recent_data[starttime:]
        logging.debug(
            "datastream.recent_data after drop:" + pp(self.recent_data, output=False)
        )

        # return as ndarray
        return avg.values

class DatastreamManager:
    """Stores and updates active SLM Datastreams."""

    def __init__(self, store_url, mqtt_client, filter_mics, auth=None):
        self.st_store = SensorThingStore(store_url, auth=auth)
        self.datastreams = {}
        self.mqtt_client = mqtt_client
        self.filter_mics = filter_mics

    async def update_datastreams(self):
        """Query for available SLM datastreams and subscribe."""
        logging.info("Updating list of datastreams.")

        # query for ObservedProperty CPBLZeq
        observed_property = self.st_store.get_by_query(
            "ObservedProperties", filter="equals(name, 'CPBLZeq')"
        )[0]

        # query for all Datastreams that observe above property
        all_datastreams = self.st_store.get_by_parent(
            kind="Datastreams",
            parent="ObservedProperties",
            id=observed_property["@iot.id"],
            filter=self.filter_mics,
        )
        if not all_datastreams:
            logging.warning(
                f"Could not find any datastreams with filter {self.filter_mics}."
            )
            return

        # populate dictionary with new datastreams
        polygon = Polygon(np.array(ast.literal_eval(os.environ["AREA_POLYGON"])))
        for jso in all_datastreams:
            stream = Datastream(jso)
            if stream.id not in self.datastreams:

                # only include streams with location in area_polygon
                if polygon.contains(Point(stream.location)):
                    self.datastreams[stream.id] = stream
                else:
                    logging.info(
                        f"Excluding datastream {stream.id} with name {stream.name} at {stream.location}: not in polygon."
                    )

        if not self.datastreams:
            logging.info("No datastreams matching criteria found.")
            return

        # subscribe to all datastreams
        topics = [(ds.mqtt_topic, QOS_1) for (_, ds) in self.datastreams.items()]

        suback_codes = await self.mqtt_client.subscribe(topics)
        for code, topic in zip(suback_codes, topics):
            if code not in [0, 1, 2]:
                logging.error(
                    f"""SUBACK return code {code}. Could not subscribe to
                    topic: {topic}"""
                )
            else:
                logging.info(f"Subscribed to topic: {topic}")

    async def start(self, rate=UPDATE_DATASTREAMS_PERIODE):
        """Continually update datastream subscriptions."""
        while True:
            await self.update_datastreams()
            logging.debug("Datastreams: {}".format(pp(self.datastreams, output=False)))
            await asyncio.sleep(rate)

class DataCollector:
    """Collects MQTT messages and their data."""

    def __init__(self, mqtt_client, datastreams):
        self.mqtt_client = mqtt_client
        self.datastreams = datastreams

    def process(self, message):
        """Add contents of message to datastore."""
        # keep track of datastreams using their iot.id
        datastream_id = int(re.search(r"\d+", message.topic).group())
        logging.debug(
            f"Processing message from Datastream({datastream_id})"
            + f" from topic {message.topic}"
        )

        data = CPBLZeq_Observation_to_dataframe(message)

        if datastream_id in self.datastreams:
            self.datastreams[datastream_id].add(data)
        else:
            logging.error(
                f"Received message from Datastream({datastream_id}), altough not subscribed ..."
            )

        logging.debug(
            "Added data to datastreams:"
            + pp(self.datastreams[datastream_id], output=False)
        )

    async def start(self):
        """Listen to SLM observations and process incoming messages."""

        timeout = 60

        while True:

            try:
                message = await self.mqtt_client.deliver_message(timeout=timeout)
                logging.info(f"Received message on {message.topic}")
                logging.debug(f"Payload: {pp(json.loads(message.data), output=False)}")

                self.process(message)

            except ClientException as ce:
                logging.error("Client exception: %s" % ce)
            except asyncio.TimeoutError:
                logging.error(
                    f"Timeout error: did not receive message since {timeout}s"
                )

class MapMaker:
    def __init__(self, datastreams, mqtt_client, mqtt_shm_topics):
        self.datastreams = datastreams
        self.mqtt_client = mqtt_client
        self.mqtt_shm_topics = mqtt_shm_topics

    async def make_map(self):
        now = datetime.utcnow()
        starttime = now - timedelta(
            seconds=MAP_LATENCY + NEW_MAP_PERIODE + MAP_AVERAGING_TIME
        )
        endtime = now - timedelta(seconds=MAP_LATENCY)

        logging.debug(f"starttime: {starttime}, endtime: {endtime}")

        spl_values = []
        slm_latlon = []
        for (_, ds) in self.datastreams.items():
            avg = ds.average_and_drop_old_data(starttime, endtime)
            if avg is not None:
                # there was data to average in that microphone
                spl_values.append(avg)
                slm_latlon.append(ds.location)

        if not spl_values:
            logging.info("Couldn't create map: no data.")
            return None

        logging.info("Making a map.")

        slm_latlon = np.array(slm_latlon).T  # to shape 2 x Nslm
        spl_values = np.array(spl_values).T  # to shape Nf x Nslm

        logging.debug(
            pp({"Locations": slm_latlon, "averages": spl_values}, output=False)
        )

        temperature = float(os.environ["TEMPERATURE"])
        s_latlon = np.array(
            ast.literal_eval(os.environ["SOURCES"])
        ).T  # to shape 2 x Ns
        walls = os.environ["WALLS"]
        if walls:
            wall_latlon = np.array(
                ast.literal_eval(os.environ["WALLS"])
            ).T  # to shape 2 x 2 x Nw
        else:
            wall_latlon = None
        polygon = np.array(ast.literal_eval(os.environ["AREA_POLYGON"])).T
        cellsize = float(os.environ["CELLSIZE"])
        source_direction = np.array(ast.literal_eval(os.environ["SOURCES_DIRECTION"]))

        assert s_latlon.ndim == 2
        assert s_latlon.shape[0] == 2
        if walls:
            assert wall_latlon.ndim == 3
            assert wall_latlon.shape[0] == 2 and wall_latlon.shape[1] == 2
        assert polygon.ndim == 2
        assert polygon.shape[0] == 2
        assert isinstance(cellsize, float)
        assert FREQUENCIES.size == spl_values.shape[0]
        assert isinstance(temperature, float) or isinstance(temperature, int)

        # make a grid of points that includes the polygon
        r_latlon_all, mask_inside_polygon, grid_shape = grid_area(
            polygon, dx=cellsize, dy=cellsize
        )
        r_latlon_inside_polygon = r_latlon_all[:, mask_inside_polygon]

        # inside the polygon, one has the computed values
        try:
            shm = calculate_shm_dir(
                s_latlon=s_latlon,
                slm_latlon=slm_latlon,
                wall_latlon=wall_latlon,
                r_latlon=r_latlon_inside_polygon,
                Lp=spl_values,
                f=FREQUENCIES,
                T=temperature,
                alpha=source_direction,
            )
        except ValueError as e:
            logging.error(f"Could not create map: {e}")
            return

        for weighting, topic in zip(WEIGHTINGS, self.mqtt_shm_topics):

            # apply weightings
            wfunc = weightfunc(weighting[0])
            shm_weighted = np.round(shm, 1) + wfunc(FREQUENCIES.astype(float))[:, None]

            # sum or not
            if weighting.endswith("fullband"):
                # compute a fullband spectrum
                band_frequencies = ["fullband"]

                # outside the polygon everything is NaN
                L_all = np.empty((len(band_frequencies), r_latlon_all.shape[1]), dtype=object)

                # sum over rms pressures and convert back
                rms = np.sum(10 ** (shm_weighted / 10), axis=0)[None]
                shm_weighted = 10 * np.log10(rms)
            else:
                # compute 1/3 octave wise
                band_frequencies = FREQUENCIES.tolist()
                L_all = np.empty((len(band_frequencies), r_latlon_all.shape[1]), dtype=object)

            # prepare output
            L_all.fill(None)
            L_all[:, mask_inside_polygon] = shm_weighted
            L_all_final_shape = L_all.reshape((-1, grid_shape[0], grid_shape[1]))

            msg = {
                "phenomenonTime": endtime.isoformat(sep="T", timespec="seconds") + "Z",
                "resultTime": datetime.utcnow().isoformat() + "Z",
                "result": {
                    "starttime": starttime.isoformat(sep="T", timespec="seconds"),
                    "endtime": endtime.isoformat(sep="T", timespec="seconds"),
                    "timeStamp": endtime.isoformat(sep="T", timespec="seconds"),
                    "lat_0": r_latlon_all[0, 0],
                    "lon_0": r_latlon_all[1, 0],
                    "nfreq": len(band_frequencies),
                    "nrow": L_all_final_shape.shape[1],
                    "ncols": L_all_final_shape.shape[2],
                    "cellsize": cellsize,
                    "data": L_all_final_shape.tolist(),
                    "bandFrequencies": band_frequencies,
                    "unit": f"SPL dB {weighting}",
                    "input_positions": {
                        "sources_latlon": s_latlon.tolist(),
                        "slm_latlon": slm_latlon.tolist(),
                        "walls_latlon": wall_latlon.tolist()
                        if wall_latlon is not None
                        else [],
                        "shm_area": polygon.tolist(),
                    },
                },
            }

            await self.mqtt_client.publish(topic, json.dumps(msg).encode())
            logging.info(f"Published map on {topic}")

    async def start(self):
        while True:
            await self.make_map()
            await asyncio.sleep(NEW_MAP_PERIODE)

def register_shm_as_service(catalogue_url):
    """Register service at OGC service catalogue.

    Returns MQTT topics for the three SHM variations and MQTT broker address and port
    """
    topics = []
    for weighting in WEIGHTINGS:
        msg = {
            "externalId": f"SoundHeatMap/{SHM_ID}/{weighting}",
            "metadata": "SoundHeatMapGenerator",
            "sensorType": f"SoundHeatMap/L{weighting}eq",
            "unitOfMeasurement": f"Sound Heat Map with {weighting}-weighted values.",
            "fixedLatitude": 0,
            "fixedLongitude": 0,
        }
        r = requests.post(catalogue_url + "/SearchOrCreateOGCDataStreamId", json=msg)
        r.raise_for_status()

        # get MQTT topic and server
        data = r.json()

        topic = data["mqttTopic"]
        mqtt_address, mqtt_port = data["mqttServer"].split(":")
        mqtt_address = "mqtt://" + mqtt_address

        logging.info(
            f"Registered {weighting} weighting at topic {topic}, {mqtt_address}:{mqtt_port}"
        )

        topics.append(topic)

    return topics, mqtt_address, mqtt_port

async def main():
    parser = optparse.OptionParser()
    parser.add_option("-l", "--logging-level", help="Logging level")
    parser.add_option("-f", "--logging-file", help="Logging file name")
    parser.add_option("-a", "--account", help="account")
    parser.add_option("-p", "--password", help="password")
    (options, args) = parser.parse_args()

    if options.account and options.password:
        # auth data provided in command line
        auth = (options.account, options.password)
    else:
        # auth data provided from environment variables
        auth = AUTH

    # initialize logging
    logging_levels = {
        "critical": logging.CRITICAL,
        "error": logging.ERROR,
        "warning": logging.WARNING,
        "info": logging.INFO,
        "debug": logging.DEBUG,
    }
    logging_level = logging_levels.get(options.logging_level, logging.INFO)
    logging.basicConfig(
        level=logging_level,
        filename=options.logging_file,
        format="%(asctime)s %(levelname)-6s %(name)s [%(filename)s:%(lineno)d] %(message)s",
        datefmt="%d-%m-%Y:%H:%M:%S",
    )

    logging.getLogger("pyost").addHandler(logging.StreamHandler())
    logging.getLogger("pyost").setLevel(logging_level)
    logging.getLogger(__name__).setLevel(logging_level)
    coloredlogs.install(level=logging_level, logger=logging.getLogger(__name__))

    mqtt_topics, mqtt_address, mqtt_port = register_shm_as_service(
        os.environ["CATALOG_URL"]
    )

    while True:
        try:
            # Connect to MQTT broker
            mqtt_client = MQTTClient()
            retcode = await mqtt_client.connect(mqtt_address, mqtt_port)
            if retcode == 0:
                logging.info(f"Connected to MQTT broker: {mqtt_address}:{mqtt_port}")
            else:
                raise ConnectionError(
                    f"Could not connect to broker, CONNACK code {retcode}"
                )

            # Connect to OGC SensorThing store server
            datastream_manager = DatastreamManager(
                GOST_URL, mqtt_client, FILTER_MICS, auth=auth
            )
            data_collector = DataCollector(mqtt_client, datastream_manager.datastreams)
            map_maker = MapMaker(
                datastream_manager.datastreams, mqtt_client, mqtt_topics
            )

            # run concurrently
            await asyncio.gather(
                datastream_manager.start(), data_collector.start(), map_maker.start()
            )
        except KeyboardInterrupt:
            logging.info("KeyboardInterrupt. Shutting down ...")
            sys.exit(0)
        except ImportError:
            logging.error(traceback.format_exc())
            sys.exit(1)
        except ConnectionError:
            logging.error(traceback.format_exc())
            logging.error("Restarting...")
        except Exception:
            logging.error(traceback.format_exc())
            logging.error("Restarting...")
        finally:
            mqtt_client.disconnect()
            break

if __name__ == "__main__":
    asyncio.run(main())


"""
The pycity_scheduling framework

Copyright (C) 2022,
Institute for Automation of Complex Power Systems (ACS),
E.ON Energy Research Center (E.ON ERC),
RWTH Aachen University

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit
persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""

import numpy as np
import pyomo.environ as pyomo

from pycity_scheduling.classes.thermal_entity_cooling import ThermalEntityCooling
from pycity_scheduling.classes.thermal_entity_heating import ThermalEntityHeating
from pycity_scheduling.classes.electrical_entity import ElectricalEntity

class EntityContainer(ThermalEntityCooling, ThermalEntityHeating, ElectricalEntity):
    """
    Base class for entities containing other entities.

    `p_th` and `p_el` imbalances are propagated to this entities variables.
    During calls to its scheduling functions, the contained entities are also
    called with the same parameters.

    Notes
    -----
    - EntityContainers offer sets of constraints for operation. The following
      constraints are added.

    .. math::
        p_{th\\_cool} &=& \\sum_i p_{th\\_cool\\_i} \\\\
        p_{th\\_heat} &=& \\sum_i p_{th\\_heat\\_i} \\\\
        p_{el} &=& \\sum_i p_{el\\_i}

    - :math:`p_{th\\_cool\\_i}`, :math:`p_{th\\_heat\\_i}`, and :math:`p_{el\\_i}` are the variables from lower
      entities. The Bounds from TEC, TEH, and EE are removed.
    """

    def populate_model(self, model, mode="convex"):
        """
        Add entity block and lower entities blocks to pyomo ConcreteModel.

        Call both parent's `populate_model` methods and set variables lower
        bounds to `None`. Then call `populate_model` method of all contained
        entities and add constraints that the sum of their variables for each
        period equals the corresponding own variable.

        Parameters
        ----------
        model : pyomo.ConcreteModel
        mode : str, optional
            Specifies which set of constraints to use.

            - `convex`  : Use linear constraints
            - `integer`  : Use same constraints as convex mode
        """
        super().populate_model(model, mode)
        m = self.model

        if mode in ["convex", "integer"]:
            p_th_cool_var_list = []
            p_th_heat_var_list = []
            p_el_var_list = []
            for entity in self.get_lower_entities():
                entity.populate_model(model, mode)
                if isinstance(entity, ThermalEntityCooling):
                    p_th_cool_var_list.append(entity.model.p_th_cool_vars)
                if isinstance(entity, ThermalEntityHeating):
                    p_th_heat_var_list.append(entity.model.p_th_heat_vars)
                if isinstance(entity, ElectricalEntity):
                    p_el_var_list.append(entity.model.p_el_vars)

            m.p_th_cool_vars.setlb(None)
            m.p_th_heat_vars.setlb(None)
            m.p_el_vars.setlb(None)

            def p_th_cool_sum_rule(model, t):
                return model.p_th_cool_vars[t] == pyomo.quicksum(p_th_Cool_var[t] for
                                                                 p_th_Cool_var in p_th_cool_var_list)
            m.p_th_cool_constr = pyomo.Constraint(m.t, rule=p_th_cool_sum_rule)

            def p_th_heat_sum_rule(model, t):
                return model.p_th_heat_vars[t] == pyomo.quicksum(p_th_Heat_var[t] for
                                                                 p_th_Heat_var in p_th_heat_var_list)
            m.p_th_heat_constr = pyomo.Constraint(m.t, rule=p_th_heat_sum_rule)

            def p_el_sum_rule(model, t):
                return model.p_el_vars[t] == pyomo.quicksum(p_el_var[t] for p_el_var in p_el_var_list)
            m.p_el_constr = pyomo.Constraint(m.t, rule=p_el_sum_rule)
        else:
            raise ValueError(
                "Mode %s is not implemented by class EntityContainer." % str(mode)
            )
        return

    def update_model(self, mode=""):
        super().update_model(mode)
        for entity in self.get_lower_entities():
            entity.update_model(mode)
        return

    def update_schedule(self):
        super().update_schedule()
        for entity in self.get_lower_entities():
            entity.update_schedule()
        return

    def reset(self, schedule=None):
        super().reset(schedule)
        for entity in self.get_lower_entities():
            entity.reset(schedule)
        return

    def get_lower_entities(self):
        raise NotImplementedError


# Import smorgasbord
import os
import sys
import gc
import shutil
import pdb
import time
import copy
import numbers
import warnings
import psutil
import random
import multiprocessing as mp
import numpy as np
import scipy.optimize
import astropy.io.fits
import astropy.wcs
import astropy.convolution
import ChrisFuncs
import ChrisFuncs.Photom
import ChrisFuncs.FromGitHub
import CAAPR

# The aperture-fitting sub-pipeline
def SubpipelineAperture(source_dict, band_dict, kwargs_dict):
    source_id = source_dict['name']+'_'+band_dict['band_name']

    # Carry out small random wait, to stop RAM checks from syncing up later
    time.sleep(5.0*np.random.rand())

    # Perform initial checks of target file type and location; return if not present
    in_fitspath, file_found = CAAPR.CAAPR_Pipeline.FilePrelim(source_dict, band_dict, kwargs_dict)
    if file_found == False:
        return None

    # Create the pod (Photometry Organisation Dictionary), which will read in the FITS file, and bundle all the photometry data for this source & band into one dictionary to be passed between functions
    pod = CAAPR.CAAPR_Pipeline.PodInitiate(in_fitspath, source_dict, band_dict, kwargs_dict)

    # Run pod through preliminary processing, to determine initial quantities; if target not within bounds of map, end processing here
    pod = CAAPR.CAAPR_Pipeline.MapPrelim(pod, source_dict, band_dict)
    if pod['within_bounds']==False:
        return None
    CAAPR.CAAPR_IO.MemCheck(pod)

    # Check if this band is to be excluded from aperture-fitting; if so, return null aperture information
    pod = ExcludeAperture(pod, source_dict, band_dict, kwargs_dict)
    if pod['band_exclude']==True:
        return pod['null_output_dict']

    # If star-removal is required, run pod through AstroMagic
    pod = CAAPR.CAAPR_AstroMagic.Magic(pod, source_dict, band_dict, kwargs_dict)

    # Run pod through function that determines aperture shape, to provide preliminary estimate to facilitate removal of large-scale sky
    pod = ApertureShape(pod)

    # Run pod through function that removes large-scale sky using a 2-dimensional polynomial filter
    pod = CAAPR.CAAPR_Pipeline.PolySub( pod, 2.0*pod['semimaj_initial_pix'], pod['opt_axial_ratio'], pod['opt_angle'], instant_quit=max([not kwargs_dict['polysub'],pod['band_exclude']]) )

    # If sky polynomial removed, run pod through function that determines aperture shape, to provide final estiamte
    if pod['sky_poly']!=False:
        pod = ApertureShape(pod)

    # Run pod through function that determines aperture size
    pod = ApertureSize(pod, band_dict)

    # If thumbnail images have been requested, save a copy of the current image (ie, with any star and/or background subtaction)
    if kwargs_dict['thumbnails']==True:
        astropy.io.fits.writeto(os.path.join(kwargs_dict['temp_dir_path'],'Processed_Maps',source_id+'.fits'), pod['cutout'], header=pod['in_header'], overwrite=True)

    # Now return final aperture informaton to main pipeline, and clean up garbage
    output_dict = {'band_name':band_dict['band_name'],
                   'opt_semimaj_arcsec':pod['opt_semimaj_arcsec'],
                   'opt_axial_ratio':pod['opt_axial_ratio'],
                   'opt_angle':pod['opt_angle']}
    gc.collect()
    del(pod)
    return output_dict

# Define function that determines the shape (not the size) of the source aperture in this band
def ApertureShape(pod):
    if pod['band_exclude']==True:
        return pod
    verbose = pod['verbose']
    if pod['verbose']: print '['+pod['id']+'] Commencing determination of appropriate axial ratio and positional angle for source aperture.'

    # Make preliminary per-pixel noise measurement by iteratively sigma-clipping cutout
    if verbose: print '['+pod['id']+'] Making preliminary per-pixel noise measurement.'
    clip_value = ChrisFuncs.SigmaClip(pod['cutout'], tolerance=0.001, sigma_thresh=3.0, median=True)
    noise_value = clip_value[0]
    field_value = clip_value[1]

    # Find all significant pixels that are connected to the region of the source (ie, withing a beam-width of the provided target coords)
    if verbose: print '['+pod['id']+'] Finding contiguous significant pixels around target.'
    semimaj_initial = int(round(pod['beam_pix']*1.0))
    cutoff = field_value + (4.0*noise_value)
    #cont_structure = np.array([[1]*(2*semimaj_initial)]*(2*semimaj_initial))
    cont_array_prelim = ChrisFuncs.Photom.ContiguousPixels(pod['cutout'], semimaj_initial, pod['centre_i'], pod['centre_j'], cutoff)#, custom_structure=cont_structure)

    # Use binary erosion to remove thin artefacts (primarily diffraction spikes)
    erode_size = int(np.ceil(3.0*pod['beam_pix']))
    erode_centre = (0.5*float(erode_size))-0.5
    erode_structure = ChrisFuncs.Photom.EllipseMask(np.zeros([erode_size,erode_size]), pod['beam_pix'], 1.0, 0.0, erode_centre, erode_centre)
    erode_array = scipy.ndimage.morphology.binary_erosion(cont_array_prelim, structure=erode_structure).astype(int)
    cont_array = ChrisFuncs.Photom.ContiguousPixels(erode_array, semimaj_initial, pod['centre_i'], pod['centre_j'], 1E-50)

    # If remainging contiguous pixel region has same or fewer number of pixels than erosion structure, replace with erosion sturcture
    if np.sum(cont_array)<=(np.sum(erode_structure)):
        cont_array = erode_structure

    # Find ellipse that best fits outline of contiguous region
    if verbose: print '['+pod['id']+'] Fitting ellipse to perimeter of contiguous significant pixels.'
    cont_x = ((np.where(cont_array==1))[1])
    cont_y = ((np.where(cont_array==1))[0])
    if cont_x.shape[0]>10:
        try:
            cont_ellipse = ChrisFuncs.Photom.EllipseFit(cont_x, cont_y)
            opt_axial_ratio = max(cont_ellipse[1]) / min(cont_ellipse[1])
            opt_angle = cont_ellipse[2]
            semimaj_initial = max(cont_ellipse[1])
        except:
            opt_axial_ratio = 1.0
            opt_angle = 0.0

    # If too few significant pixels, default to circular aperture
    else:
        opt_axial_ratio = 1.0
        opt_angle = 0.0
    if verbose: print '['+pod['id']+'] Ellipse angle: '+str(ChrisFuncs.FromGitHub.randlet.ToPrecision(opt_angle,4))+' degrees; Ellipse axial ratio: '+str(ChrisFuncs.FromGitHub.randlet.ToPrecision(opt_axial_ratio,4))+'.'

    # Clean garbage, record results to pod, and return
    gc.collect()
    pod['cutout_clip'] = clip_value
    pod['opt_axial_ratio'] = opt_axial_ratio
    pod['opt_angle'] = opt_angle
    pod['semimaj_initial_pix'] = semimaj_initial
    pod['semimaj_initial_arcsec'] = semimaj_initial * pod['pix_arcsec']
    return pod

# Define function that determines the size of the source aperture in this band
def ApertureSize(pod, band_dict):
    if pod['band_exclude']:
        return pod
    if pod['verbose']: print '['+pod['id']+'] Commencing determination of appropriate size for source aperture.'
    verbose = pod['verbose']

    # Define sub-function that determines SNR of a defined annulus; and if requested, determines residual between the SNR of a defined annulus, and a target SNR of 2
    def AnnulusSNR(semimaj, pod, cutout, width, i_trans, j_trans, residual):
        semimaj = semimaj[0]
        sig_annulus = ChrisFuncs.Photom.AnnulusQuickSum(cutout, semimaj, width, pod['opt_axial_ratio'], pod['opt_angle'], pod['centre_i'], pod['centre_j'], i_trans, j_trans)
        sig_value = ChrisFuncs.SigmaClip(sig_annulus[2], tolerance=0.005, sigma_thresh=2.0, median=True)[1]#np.median(sig_annulus[2])
        noise_value = pod['cutout_clip'][0]
        field_value = pod['cutout_clip'][1]
        ann_SNR = (sig_value - field_value) / noise_value
        if residual==True:
            ann_residual = abs(2.0-ann_SNR)
            #print 'SNR: '+str(ann_SNR)+', Semi-Maj: '+str(semimaj)+', Residual:'+str(ann_residual)
            return ann_residual
        elif residual==False:
            #print 'SNR: '+str(ann_SNR)+', Semi-Maj: '+str(semimaj)
            return ann_SNR

    # Construct kernel with FWHM equal to 3 beam-widths, by which to smooth map
    if verbose: print '['+pod['id']+'] Convolving map to lower resolution (twice the beam width) for radial analysis.'
    pix_size = pod['pix_arcsec']
    res_in = band_dict['beam_arcsec']
    res_out = 2.0*band_dict['beam_arcsec']#36.0
    kernel_fwhm = np.sqrt( (res_out/pix_size)**2.0 - (res_in/pix_size)**2.0 )

    # Determine if map contains NaN pixels, excluding those that simply represent edge of map
    cutout_prelabel = pod['cutout'].copy()
    cutout_prelabel = cutout_prelabel.byteswap().newbyteorder().astype('float64')
    cutout_prelabel[ np.where(np.isnan(cutout_prelabel)) ] = 0.0
    cutout_label = scipy.ndimage.label(cutout_prelabel)
    cutout_preconv = pod['cutout'].copy()
    cutout_preconv[ np.where(cutout_label==0) ] = 0.0

    # If map contains no NaN pixels, smooth it the quick Scipy way
    cutout_unconv = pod['cutout'].copy()
    if np.where(np.isnan(cutout_preconv)==True)[0].shape[0]==0:
        if verbose: print '['+pod['id']+'] No NaN pixels within coverage area; convolving using quick method.'
        pod['cutout'] = scipy.ndimage.filters.gaussian_filter(cutout_preconv, kernel_fwhm)

    # Else if map contains NaNs, do it the robust (but very-very slow, very-very memory intensive) Astropy way
    else:
        if verbose: print '['+pod['id']+'] NaN pixels within coverage area; convolving using slower NaN-compatible method.'
        CAAPR.CAAPR_IO.MemCheck(pod, thresh_factor=20.0)
        kernel = astropy.convolution.kernels.Gaussian2DKernel(kernel_fwhm)
        pod['cutout'] = astropy.convolution.convolve_fft(pod['cutout'], kernel, nan_treatment='interpolate', normalize_kernel=True, allow_huge=True)
        pod['cutout'][ np.where( np.isnan(cutout_unconv)==True ) ] = np.NaN

    # Prepare arrays of transposed coordinates, to allow for rapid radial evaluating
    if verbose: print '['+pod['id']+'] Constructing arrays of transposed radial coordinates.'
    coords_trans = ChrisFuncs.Photom.AnnulusQuickPrepare(pod['cutout'], pod['opt_angle'], pod['centre_i'], pod['centre_j'])
    i_trans, j_trans = coords_trans[0], coords_trans[1]

    # To start with, to make new estimate of map noise that isn't contaminated by the target galaxy, by masking all pixels beyond semi-major axis suggested by contiguous significant pixels.
    brute_mask = ChrisFuncs.Photom.EllipseMask(pod['cutout'], pod['semimaj_initial_pix'], pod['opt_axial_ratio'], pod['opt_angle'], pod['centre_i'], pod['centre_j'])
    cutout_brute_masked = pod['cutout'].copy()
    cutout_brute_masked[ np.where( brute_mask==1 ) ] = np.nan
    cutout_clip_masked = ChrisFuncs.SigmaClip(cutout_brute_masked, tolerance=0.001, sigma_thresh=3.0, median=True)
    pod['cutout_clip'] = cutout_clip_masked

    # Now, perform a coarse brute force ckeck of a small number of radii over a wide range, to find rough location of edge of the source
    if verbose: print '['+pod['id']+'] Finding size of target source with coarse analysis.'
    ann_brute_range = np.linspace(0.75*pod['semimaj_initial_pix'], 2.25*pod['semimaj_initial_pix'], num=15)
    ann_brute_range = ann_brute_range[::-1]
    ann_brute_width = abs( ann_brute_range[1] - ann_brute_range[0] )
    snr_success = False
    for i in range(0, len(ann_brute_range)):
        ann_brute_snr = AnnulusSNR([ann_brute_range[i]], pod, pod['cutout'], ann_brute_width, i_trans, j_trans, False)
        if ann_brute_snr>2:
            snr_success = True
            #ann_brute_semimaj = ann_brute_range[i-1]
            ann_bounds = [( ann_brute_range[ max(i-2,0) ], ann_brute_range[ min(i+1,len(ann_brute_range)-1) ] )]
            if verbose: print '['+pod['id']+'] Course analysis finds that radial SNR=2 between semi-major axes of '+str(ChrisFuncs.FromGitHub.randlet.ToPrecision(ann_bounds[0][1]*pod['pix_arcsec'],4))+' and '+str(ChrisFuncs.FromGitHub.randlet.ToPrecision(ann_bounds[0][0]*pod['pix_arcsec'],4))+' arcseconds.'
            break

    # If SNR=2 threshold not reached, set to default minimum semi-major axis
    if snr_success==False:
        ann_bounds = [( ann_brute_range[i], np.floor(pod['cutout'].shape[0]/2.0) )]
        opt_semimaj_pix = pod['beam_pix'] * 0.0
        opt_semimaj_arcsec = opt_semimaj_pix * pod['pix_arcsec']
        if verbose: print '['+pod['id']+'] No SNR=2 threshold found; hence reverting to default minimum aperture size of one beam-width.' #+str(ChrisFuncs.FromGitHub.randlet.ToPrecision(opt_semimaj_arcsec,4))+' arcseconds.'
    else:

        # Now use scipy differential evolution optimisation to find, with more precision, the semi-major axis at which annulus falls to a SNR of 2
        ann_beams = 1.0
        ann_width = np.ceil( ann_beams * pod['beam_pix'] )
        if verbose: print '['+pod['id']+'] Refining size of target source with more precise analysis.'
        ann_fit = scipy.optimize.differential_evolution(AnnulusSNR, ann_bounds, args=(pod, pod['cutout'], ann_width, i_trans, j_trans, True), maxiter=5, popsize=10, polish=False)
        """
        ann_guess = ann_brute_semimaj
        ann_fit = scipy.optimize.minimize(AnnulusSNR, ann_guess, args=(pod, pod['cutout'], ann_width, i_trans, j_trans))#method='Nelder-Mead', tol=1E-4, method='L-BFGS-B', bounds=[(pod['semimaj_initial_pix'], None)],
        minimizer_kwargs = {'args':(pod, pod['cutout'], ann_width, i_trans, j_trans)}
        ann_fit = scipy.optimize.basinhopping(AnnulusSNR, ann_guess, T=0.1, stepsize=5.0, minimizer_kwargs=minimizer_kwargs)
        """
        # Extract results from fitting
        opt_semimaj_pix = ann_fit['x'][0]
        opt_semimaj_arcsec = opt_semimaj_pix * pod['pix_arcsec']
        if verbose: print '['+pod['id']+'] Precision analysis finds that radial SNR=2 at semi-major axis of '+str(ChrisFuncs.FromGitHub.randlet.ToPrecision(opt_semimaj_arcsec,4))+' arcseconds.'
        """
        # For small sources, default to minimum semi-major axis of one beam-width
        if opt_semimaj_pix<(ann_beams*2.0):
            opt_semimaj_pix = pod['beam_pix'] * 2.0
            opt_semimaj_arcsec = opt_semimaj_pix * pod['pix_arcsec']
            if verbose: print '['+pod['id']+'] Semi-major axis at which SNR=2 is less than one beam-width; hence reverting to two beam-width minimum value of '+str(ChrisFuncs.FromGitHub.randlet.ToPrecision(opt_semimaj_arcsec,4))+' arcseconds.'
        """
    # Establish what fraction of the pixels inside a band's aperture are NaNs
    pix_good = ChrisFuncs.Photom.EllipseQuickSum(pod['cutout'], opt_semimaj_pix, pod['opt_axial_ratio'], pod['opt_angle'], pod['centre_i'], pod['centre_j'], i_trans, j_trans)[1]
    pix_tot = np.where( ChrisFuncs.Photom.EllipseMask(pod['cutout'], opt_semimaj_pix, pod['opt_axial_ratio'], pod['opt_angle'], pod['centre_i'], pod['centre_j']) == 1 )[0].shape[0]
    if pix_tot==0.0:
        pix_good_frac = 0.0
    else:
        pix_good_frac = float(pix_good) / float(pix_tot)

    # Before final reporting, tidy up and return to unconvolved cutout
    pod['cutout'] = cutout_unconv
    del(cutout_unconv)
    gc.collect()

    # If more than 10% of the pixels in the aperture are NaNs, report default values for aperture dimensions
    if pix_good_frac<0.9:
        opt_semimaj_pix = pod['beam_pix'] * 1.0
        opt_semimaj_arcsec = opt_semimaj_pix * pod['pix_arcsec']
        if verbose: print '['+pod['id']+'] More than 10% of pixels in fitted aperture are NaNs; hence reverting to default minimum aperture size of one beam-width.' #+str(opt_semimaj_arcsec)[:7]+' arcseconds.'
        pod['opt_semimaj_arcsec'] = opt_semimaj_arcsec
        pod['opt_axial_ratio'] = 1.0
        pod['opt_angle'] = 0.0

    # If dimensions are otherwise default, report to pod
    if abs( opt_semimaj_arcsec**2.0 - (0.5*band_dict['beam_arcsec'])**2.0 )**0.5<=0.0:
        opt_semimaj_pix = pod['beam_pix'] * 1.0
        pod['opt_semimaj_arcsec'] = opt_semimaj_pix * pod['pix_arcsec']
        pod['opt_axial_ratio'] = 1.0
        pod['opt_angle'] = 0.0

    # Else deconvolve aperture semi-major axis with beam, by subtracting in quadrature, and record dimensions to pod
    else:
        adj_semimaj_arcsec = 0.5 * np.abs( (2.0*opt_semimaj_arcsec)**2.0 - band_dict['beam_arcsec']**2.0 )**0.5
        opt_semimin_arcsec = opt_semimaj_arcsec / pod['opt_axial_ratio']
        adj_semimin_arcsec = 0.5 * np.abs( (2.0*opt_semimin_arcsec)**2.0 - band_dict['beam_arcsec']**2.0 )**0.5
        adj_ax_ratio = adj_semimaj_arcsec / adj_semimin_arcsec
        pod['opt_semimaj_arcsec'] = adj_semimaj_arcsec
        pod['opt_semimaj_pix'] = adj_semimaj_arcsec / pod['pix_arcsec']
        pod['opt_axial_ratio'] = adj_ax_ratio

    # Clean up, then return results
    gc.collect()
    return pod

# Define function that combines a set of apertures for a given source into
def CombineAperture(aperture_output_list, source_dict, kwargs_dict):
    if kwargs_dict['verbose']: print '['+source_dict['name']+'] Combining individual apertures from all bands to generate final aperture.'

    # Extract various aperture values
    semimaj_arcsec_list = []
    axial_ratio_list = []
    angle_list = []
    for aperture in aperture_output_list:
        if aperture==False:
            continue
        try:
            semimaj_arcsec_list.append( aperture['opt_semimaj_arcsec'] )
            axial_ratio_list.append( aperture['opt_axial_ratio'] )
            angle_list.append( aperture['opt_angle'] )
        except:
            pdb.set_trace()

    # Check to see if any bands have been designated for aperture consideration
    if np.nanmax(np.array(axial_ratio_list)) == 0.0:
        raise Exception('No usable apertures found; probably because no bands have consider_aperture set to True in the bands table')

    # Find largest semi-major axis, and use to define size of enclosisity array (which will have pixels some fraction the size of the smallest semi-major axis)
    semimaj_max = np.nanmax(semimaj_arcsec_list) #semimaj_min = np.nanmin(semimaj_arcsec_list)
    ap_array_pix_size = 0.005 * semimaj_max
    ap_array_scale = int( np.round( semimaj_max / ap_array_pix_size ) )
    ap_array = np.zeros([ int(1+(2.2*ap_array_scale)), int(1+(2.2*ap_array_scale)) ])
    centre_i, centre_j = 1+(1.1*ap_array_scale), 1+(1.1*ap_array_scale)
    semimaj_pix_list = np.array(semimaj_arcsec_list) / ap_array_pix_size

    # Loop over each aperture, adding to enclosisity array
    for a in range(0, len(semimaj_pix_list)):
        if np.isnan(semimaj_pix_list[a])==False:
            ap_mask = ChrisFuncs.EllipseMask(ap_array, semimaj_pix_list[a], axial_ratio_list[a], angle_list[a], centre_i, centre_j)
            ap_array[ np.where( ap_mask==1 ) ] += 1
    #ChrisFuncs.Cutout(ap_array, '/home/saruman/spx7cjc/DustPedia/Ap.fits')

    # Find ellipse that traces edge of enclosisity region
    cont_rad_initial_pix = 2.0#( semimaj_min / np.nanmax(axial_ratio_list) ) / ap_array_pix_size
    cont_array = ChrisFuncs.Photom.ContiguousPixels(ap_array, cont_rad_initial_pix, centre_i, centre_j, 0.1)
    cont_x = ((np.where(cont_array==1))[1])
    cont_y = ((np.where(cont_array==1))[0])
    if cont_x.shape[0]>10:
        try:
            cont_ellipse = ChrisFuncs.Photom.EllipseFit(cont_x, cont_y)
            cont_axial_ratio = max(cont_ellipse[1]) / min(cont_ellipse[1])
            cont_angle = cont_ellipse[2]
            cont_semimaj_pix = max([ cont_ellipse[1].max(), np.nanmax(semimaj_pix_list) ])
        except:
            cont_axial_ratio = 1.0
            cont_angle = 0.0
            cont_semimaj_pix = max([ cont_ellipse[1].max(), np.nanmax(semimaj_pix_list) ])
    else:
        pdb.set_trace()
        cont_axial_ratio = 1.0
        cont_angle = 0.0
        cont_semimaj_pix = 2.0 * np.max('beam_width')

    # Convert final semi-major axis back to arcsec and apply expanson factor, then clean garbage and return results
    if isinstance(kwargs_dict['expansion_factor'], float) or isinstance(kwargs_dict['expansion_factor'], int):
        expansion_facor = float(kwargs_dict['expansion_factor'])
    else:
        expansion_facor = 1.0
    cont_semimaj_arcsec = cont_semimaj_pix * ap_array_pix_size * expansion_facor

    # If final aperture is smaller than defined minimum aperture, switch to defined minimum
    if cont_semimaj_arcsec<source_dict['fitting_min_semimaj_arcsec']:
        if kwargs_dict['verbose']: print '['+source_dict['name']+'] Fitted aperture is smaller than minimum permitted aperture size; reverting to minimum permitted aperture size.'
        cont_semimaj_arcsec = source_dict['fitting_min_semimaj_arcsec']
        cont_axial_ratio = 1.0
        cont_angle = 0.0

    # Clean garbage and return results
    if kwargs_dict['verbose']: print '['+source_dict['name']+'] Final ellipse semi-major axis: '+str(ChrisFuncs.FromGitHub.randlet.ToPrecision(cont_semimaj_arcsec,4))+' arcsec; final ellipse angle: '+str(ChrisFuncs.FromGitHub.randlet.ToPrecision(cont_angle,4))+' '+' degrees; final ellipse axial ratio: '+str(ChrisFuncs.FromGitHub.randlet.ToPrecision(cont_axial_ratio,4))+'.'
    gc.collect()
    return [cont_semimaj_arcsec, cont_axial_ratio, cont_angle, ap_array]

# Define function to check that all bands which were supposed to undergo photometry have done so
def ApertureCheck(aperture_attempts, aperture_output_list, source_dict, bands_dict, kwargs_dict):

    # Compare number of bands with photometry returned to number of bands for which photometry was requested
    aperture_limit = 5
    if len(aperture_output_list)==len(bands_dict.keys()):
        aperture_attempts = 'Success'
        return aperture_attempts

    # Check how many attempts have been made so far, and proceed accordingly
    else:
        aperture_attempts += 1
        time.sleep(30.0)
        if aperture_attempts>=aperture_limit:
            print '['+source_dict['name']+'] Aperture fitting failed '+str(aperture_limit)+' times in succession; suggest debugging.'
            pdb.set_trace()
            raise Exception('Aperture fitting failed '+str(aperture_limit)+' times in succession; suggest debugging.')
        else:
            return aperture_attempts

# Define function that check is present band is to be excluded from aperture-fitting
def ExcludeAperture(pod, source_dict, band_dict, kwargs_dict):

    # Check if the aperture exclusion field actually contains characters; if so, make list of entries, and if not, record an empty list
    if isinstance(source_dict['aperture_bands_exclude'], basestring):
        aperture_bands_exclude = source_dict['aperture_bands_exclude'].split(';')
    elif source_dict['aperture_bands_exclude']==False:
        aperture_bands_exclude = []

    # If present band is to be excluded, note this fact in pod
    if (band_dict['consider_aperture']==False) or (band_dict['band_name'] in aperture_bands_exclude):
        pod['band_exclude'] = True

    # If exclusion not required, record and return
    else:
        pod['band_exclude'] = False
        return pod

    # Set generic null aperture properties
    pod['opt_axial_ratio'] = 1.0
    pod['opt_angle'] = 0.0
    pod['opt_semimaj_arcsec'] = 0.0#( (2.0*band_dict['beam_arcsec'])**2.0 - band_dict['beam_arcsec']**2.0 )**0.5
    pod['opt_semimaj_pix'] = 0.0#pod['opt_semimaj_arcsec'] / pod['pix_arcsec']
    pod['semimaj_initial_pix'] = 0.0#pod['opt_semimaj_arcsec'] / pod['pix_arcsec']

    # Create aperture output dictionry containing null values
    output_dict = {'band_name':band_dict['band_name'],
                   'opt_semimaj_arcsec':pod['opt_semimaj_arcsec'],
                   'opt_axial_ratio':pod['opt_axial_ratio'],
                   'opt_angle':pod['opt_angle']}
    pod['null_output_dict'] = output_dict

    # Return pod
    if pod['verbose']: print '['+pod['id']+'] No aperture fitting required from this source in this band.'
    return pod

# Define function that handles bands excluded from aperture fitting, so that they appear in thumbnail grid
def ExcludedThumb(source_dict, bands_dict, kwargs_dict, aperture_list, aperture_combined):

    # If thumbnails not required, end immediately
    if kwargs_dict['thumbnails']==False:
        return

    # Check if the aperture exclusion field for this source actually contains characters; if so make list of entries, else produce empty list
    if isinstance(source_dict['aperture_bands_exclude'], basestring):
        aperture_bands_exclude = source_dict['aperture_bands_exclude'].split(';')
    else:
        aperture_bands_exclude = []

    # Now consider bands which have been assigned a blancket aperture exclusion
    [ aperture_bands_exclude.append(band) for band in bands_dict.keys() if bands_dict[band]['consider_aperture']==False ]
    aperture_bands_exclude = list( set( aperture_bands_exclude ) )
    aperture_bands_exclude = np.array(aperture_bands_exclude)[ np.in1d( aperture_bands_exclude, bands_dict.keys() ) ]

    # If no bands require processing here, end immediately; else prepare to loop over bands that do require processing
    if len(aperture_bands_exclude)==0:
        return
    else:
        if kwargs_dict['verbose']: print '['+source_dict['name']+'] Preparing thumbnail data for bands excluded from aperture-fitting.'
        random.shuffle(aperture_bands_exclude)

    # Find largest beam size and outer annulus size, and hence work out thumbnail size that will contain the largest beam-convolved aperture
    beam_arcsec_max = 0.0
    outer_annulus_max = 0.0
    pix_arcsec_max = 0.0
    for band_name in bands_dict:
        in_fitspath, file_found = CAAPR.CAAPR_Pipeline.FilePrelim(source_dict, bands_dict[band_name], kwargs_dict)
        if file_found!=True:
            continue
        band_pix_matrix = astropy.wcs.WCS(astropy.io.fits.getheader(in_fitspath)).pixel_scale_matrix
        band_pix_arcsec = 3600.0 * np.sqrt( np.min(np.abs(band_pix_matrix))**2.0 + np.max(np.abs(band_pix_matrix))**2.0 )
        if band_pix_arcsec>pix_arcsec_max:
            pix_arcsec_max = band_pix_arcsec
        if bands_dict[band_name]['beam_arcsec']>beam_arcsec_max:
            beam_arcsec_max = bands_dict[band_name]['beam_arcsec']
        if bands_dict[band_name]['annulus_outer']>outer_annulus_max:
            outer_annulus_max = bands_dict[band_name]['annulus_outer']
    thumb_rad_arcsec = np.ceil( 1.0 * pix_arcsec_max ) + np.ceil( 1.75 * 0.5 * np.sqrt( (outer_annulus_max*2.0*aperture_combined[0])**2.0 + (beam_arcsec_max)**2.0 ) )
    source_dict['thumb_rad_arcsec'] = thumb_rad_arcsec

    # In standard operation, process multiple sources in parallel
    if kwargs_dict['parallel']==True:
        ex_ap_pool = mp.Pool(processes=kwargs_dict['n_proc'])
        for band in aperture_bands_exclude:
            ex_ap_pool.apply_async( ExcludedSubpipelineAperture, args=(aperture_combined, source_dict, bands_dict[band], kwargs_dict,) )
        ex_ap_pool.close()
        ex_ap_pool.join()
        del(ex_ap_pool)

    # If parallelisation is disabled, process sources one-at-a-time
    elif kwargs_dict['parallel']==False:
        for band in aperture_bands_exclude:
            ExcludedSubpipelineAperture(aperture_combined, source_dict, bands_dict[band], kwargs_dict)

# Define 'pseudo-dummy' version of the aperture sub-pipeline, to run excluded bands through
def ExcludedSubpipelineAperture(aperture_combined, source_dict, band_dict, kwargs_dict_inviolate):
    source_id = source_dict['name']+'_'+band_dict['band_name']

    # Make deep copy of kwargs dict, to disable verbosity
    kwargs_dict = copy.deepcopy(kwargs_dict_inviolate)
    kwargs_dict['verbose'] = False

    # Run through initial stages of aperture sub-pipeline, as would occur usually
    in_fitspath_prelim, file_found = CAAPR.CAAPR_Pipeline.FilePrelim(source_dict, band_dict, kwargs_dict)
    if file_found == False:
        return
    pod = CAAPR.CAAPR_Pipeline.PodInitiate(in_fitspath_prelim, source_dict, band_dict, kwargs_dict)
    pod = CAAPR.CAAPR_Pipeline.MapPrelim(pod, source_dict, band_dict)
    if pod['within_bounds']==False:
        return
    CAAPR.CAAPR_IO.MemCheck(pod)

    # Use thumbnail cutout function to create a cutout that's only as large as it needs to be for the thumbnail grid
    CAAPR.CAAPR_IO.ThumbCutout(source_dict, band_dict, kwargs_dict, pod['in_fitspath'], source_dict['thumb_rad_arcsec'])

    # Rename thumbnail cutout, and make it the 'active' map by repeating necessary processing
    thumb_output = os.path.join( kwargs_dict['temp_dir_path'], 'Processed_Maps', source_id+'_Thumbnail.fits' )
    pod['in_fitspath'] = thumb_output
    in_fitsdata = astropy.io.fits.open(pod['in_fitspath'])
    pod['in_image'] = in_fitsdata[0].data
    pod['in_header'] = in_fitsdata[0].header
    in_fitsdata.close()
    pod['in_wcs'] = astropy.wcs.WCS(pod['in_header'])
    pod['in_fitspath_size'] = float(os.stat(pod['in_fitspath']).st_size)
    thumb_centre_xy = pod['in_wcs'].wcs_world2pix( np.array([[ source_dict['ra'], source_dict['dec'] ]]), 0 )
    pod['centre_i'], pod['centre_j'] = float(thumb_centre_xy[0][1]), float(thumb_centre_xy[0][0])

    # Run thumbnail cutout thorugh AstroMagic, save result, and delete temporary files
    if kwargs_dict['starsub']==True:
        pod['cutout'] = pod['in_image'].copy()
        pod['starsub_thumbnail'] = True
        pod = CAAPR.CAAPR_AstroMagic.Magic(pod, source_dict, band_dict, kwargs_dict)
        os.remove(thumb_output)
        magic_output = os.path.join(kwargs_dict['temp_dir_path'], 'AstroMagic', band_dict['band_name'], source_dict['name']+'_'+band_dict['band_name']+'_StarSub.fits')
        if os.path.exists(magic_output):
            os.remove(magic_output)
    else:
        pod['cutout'] = pod['in_image'].copy()
        os.remove(thumb_output)

    # Save resulting cutout
    astropy.io.fits.writeto(os.path.join(kwargs_dict['temp_dir_path'],'Processed_Maps',source_id+'.fits'), pod['cutout'], header=pod['in_header'], overwrite=True)

import numpy,import scipy,import astropy


from __future__ import annotations
# dh import
try:
    from dh.udfdh import UDFDH
    from dh.dhutil import calc_batch_size, gen_batch, gen_shl_batch, tot_size, timing
    from dh.grad.rdfdh import get_H_1_ao, get_S_1_ao, generator_L_1
    from dh.grad.rdfdh import Gradients as RGradients
except ImportError:
    from pyscf.dh.udfdh import UDFDH
    from pyscf.dh.dhutil import calc_batch_size, gen_batch, gen_shl_batch, tot_size, timing
    from pyscf.dh.grad.rdfdh import get_H_1_ao, get_S_1_ao, generator_L_1
    from pyscf.dh.grad.rdfdh import Gradients as RGradients
# pyscf import
from pyscf import gto, lib, df
from pyscf.df.grad.rhf import _int3c_wrapper as int3c_wrapper
try:
    from pyscf.dftd3 import itrf
except ImportError:
    print('''Warning: dftd3 not found. You cannot using functionals with "-D3" suffix
             before installing pyscf-dftd3. See https://github.com/pyscf/dftd3 and
             https://github.com/ajz34/dh#dftd3-extension ''')
# other import
import numpy as np
import itertools
import ctypes

einsum = lib.einsum
α, β = 0, 1
αα, αβ, ββ = 0, 1, 2

@timing
def get_gradient_jk(dfobj: df.DF, C, D, D_r, Y_mo, cx, cx_n, max_memory=2000):
    mol, aux = dfobj.mol, dfobj.auxmol
    natm, nao, nmo, nocc = mol.natm, mol.nao, C.shape[-1], mol.nelec
    mocc = max(nocc)
    naux = Y_mo[0].shape[0]
    # this algorithm asserts naux = aux.nao, i.e. no linear dependency in auxiliary basis
    assert naux == aux.nao
    so = slice(0, nocc[α]), slice(0, nocc[β])

    D_r_symm = (D_r + D_r.swapaxes(-1, -2)) / 2
    D_r_ao = einsum("sup, spq, svq -> suv", C, D_r_symm, C)
    D_mo = np.zeros((2, nmo, nmo))
    for σ in (α, β):
        for i in range(nocc[σ]):
            D_mo[σ, i, i] = 1

    Y_dot_D, Y_dot_D_r = np.zeros((2, naux)), np.zeros((2, naux))
    nbatch = calc_batch_size(nmo**2, max_memory)
    for σ in (α, β):
        for i in range(nocc[σ]):
            Y_dot_D[σ] += Y_mo[σ][:, i, i]
        for saux in gen_batch(0, naux, nbatch):
            Y_dot_D_r[σ][saux] = einsum("Ppq, pq -> P", Y_mo[σ][saux], D_r_symm[σ])

    Y_ip = [np.asarray(Y_mo[σ][:, so[σ]]) for σ in (α, β)]
    L_inv, L_1_gen = generator_L_1(aux)
    int3c2e_ip1_gen = int3c_wrapper(mol, aux, "int3c2e_ip1", "s1")
    int3c2e_ip2_gen = int3c_wrapper(mol, aux, "int3c2e_ip2", "s1")
    C0 = [C[σ][:, so[σ]] for σ in (α, β)]
    D1 = [cx * D_r_symm[σ] + 0.5 * cx_n * D_mo[σ] for σ in (α, β)]
    C1 = [C[σ] @ D1[σ] for σ in (α, β)]

    grad_contrib = np.zeros((natm, 3))
    for A in range(natm):
        shA0, shA1, _, _ = mol.aoslice_by_atom()[A]
        shA0a, shA1a, _, _ = aux.aoslice_by_atom()[A]

        Y_1_mo_D_r = [np.zeros((3, naux, nocc[σ], nmo)) for σ in (α, β)]
        Y_1_dot_D, Y_1_dot_D_r = np.zeros((2, 3, naux)), np.zeros((2, 3, naux))

        pre_flop = tot_size(Y_1_mo_D_r, Y_ip, Y_1_dot_D, Y_1_dot_D_r)
        nbatch = calc_batch_size(3*(nao+mocc)*naux, max_memory, pre_flop)
        for shU0, shU1, U0, U1 in gen_shl_batch(mol, nbatch, shA0, shA1):
            su = slice(U0, U1)
            int3c2e_ip1 = int3c2e_ip1_gen((shU0, shU1, 0, mol.nbas, 0, aux.nbas))
            for σ in (α, β):
                Y_1_mo_D_r[σ] -= einsum("tuvQ, PQ, ui, vp -> tPip", int3c2e_ip1, L_inv, C0[σ][su], C1[σ])
                Y_1_mo_D_r[σ] -= einsum("tuvQ, PQ, up, vi -> tPip", int3c2e_ip1, L_inv, C1[σ][su], C0[σ])
                Y_1_dot_D[σ] -= 2 * einsum("tuvQ, PQ, uv -> tP", int3c2e_ip1, L_inv, D[σ][su])
                Y_1_dot_D_r[σ] -= 2 * einsum("tuvQ, PQ, uv -> tP", int3c2e_ip1, L_inv, D_r_ao[σ][su])

        nbatch = calc_batch_size(3*nao*(nao+mocc), max_memory, pre_flop)
        for shP0, shP1, P0, P1 in gen_shl_batch(aux, nbatch, shA0a, shA1a):
            sp = slice(P0, P1)
            int3c2e_ip2 = int3c2e_ip2_gen((0, mol.nbas, 0, mol.nbas, shP0, shP1))
            for σ in (α, β):
                Y_1_mo_D_r[σ] -= einsum("tuvQ, PQ, ui, vp -> tPip", int3c2e_ip2, L_inv[:, sp], C0[σ], C1[σ])
                Y_1_dot_D[σ] -= einsum("tuvQ, PQ, uv -> tP", int3c2e_ip2, L_inv[:, sp], D[σ])
                Y_1_dot_D_r[σ] -= einsum("tuvQ, PQ, uv -> tP", int3c2e_ip2, L_inv[:, sp], D_r_ao[σ])

        L_1 = L_1_gen(A)
        L_1_dot_inv = einsum("tRQ, PR -> tPQ", L_1, L_inv)
        for σ in (α, β):
            Y_1_mo_D_r[σ] -= einsum("Qiq, qp, tPQ -> tPip", Y_ip[σ], D1[σ], L_1_dot_inv)
            Y_1_dot_D[σ] -= einsum("Q, tPQ -> tP", Y_dot_D[σ], L_1_dot_inv)
            Y_1_dot_D_r[σ] -= einsum("Q, tPQ -> tP", Y_dot_D_r[σ], L_1_dot_inv)
            # RI-K contribution
            grad_contrib[A] += - 2 * einsum("Pip, tPip -> t", Y_ip[σ], Y_1_mo_D_r[σ])

        # RI-J contribution
        for σ, ς in itertools.product((α, β), (α, β)):
            grad_contrib[A] += (
                + einsum("P, tP -> t", Y_dot_D[σ], Y_1_dot_D_r[ς])
                + einsum("P, tP -> t", Y_dot_D_r[σ], Y_1_dot_D[ς])
                + einsum("P, tP -> t", Y_dot_D[σ], Y_1_dot_D[ς]))
    return grad_contrib

class Gradients(UDFDH, RGradients):

    def __init__(self, mol: gto.Mole, *args, skip_construct=False, **kwargs):
        if not skip_construct:
            super(Gradients, self).__init__(mol, *args, **kwargs)
        # results
        self.grad_jk = NotImplemented
        self.grad_gga = NotImplemented
        self.grad_pt2 = NotImplemented
        self.grad_enfunc = NotImplemented
        self.grad_tot = NotImplemented
        self.de = NotImplemented

    @timing
    def prepare_H_1(self):
        H_1_ao = get_H_1_ao(self.mol)
        H_1_mo = np.array([einsum("up, Auv, vq -> Apq", self.C[σ], H_1_ao, self.C[σ]) for σ in (α, β)])
        self.tensors.create("H_1_ao", H_1_ao)
        self.tensors.create("H_1_mo", H_1_mo)

    @timing
    def prepare_S_1(self):
        S_1_ao = get_S_1_ao(self.mol)
        S_1_mo = np.array([einsum("up, Auv, vq -> Apq", self.C[σ], S_1_ao, self.C[σ]) for σ in (α, β)])
        self.tensors.create("S_1_ao", S_1_ao)
        self.tensors.create("S_1_mo", S_1_mo)

    def prepare_gradient_jk(self):
        D_r = self.tensors.load("D_r")
        Y_mo = [self.tensors["Y_mo_jk" + str(σ)] for σ in (α, β)]
        # a special treatment
        cx_n = self.cx_n if self.xc_n else self.cx
        self.grad_jk = get_gradient_jk(self.df_jk, self.C, self.D, D_r, Y_mo, self.cx, cx_n, self.get_memory())

    @timing
    def prepare_gradient_gga(self):
        tensors = self.tensors
        if "rho" not in tensors:
            self.grad_gga = 0
            return self
        # --- LAZY CODE ---
        from pyscf import grad, hessian
        ni, mol, grids = self.ni, self.mol, self.grids
        natm = mol.natm
        C, D = self.C, self.D
        grad_contrib = np.zeros((natm, 3))

        xc = self.xc_n if self.xc_n else self.xc
        if self.ni._xc_type(xc) == "GGA":  # energy functional contribution
            veff_1_gga = grad.uks.get_vxc(ni, mol, grids, xc, D)[1]
            for A, (_, _, A0, A1) in enumerate(mol.aoslice_by_atom()):
                grad_contrib[A] += 2 * einsum("stuv, suv -> t", veff_1_gga[:, :, A0:A1], D[:, A0:A1])

        if self.ni._xc_type(self.xc) == "GGA":  # reference functional skeleton fock derivative contribution
            D_r = tensors.load("D_r")
            D_r_symm = (D_r + D_r.swapaxes(-1, -2)) / 2
            D_r_ao = einsum("sup, spq, svq -> suv", C, D_r_symm, C)

            F_1_ao_dfa = np.array(hessian.uks._get_vxc_deriv1(self.mf_s.Hessian(), C, self.mo_occ, 2000))
            grad_contrib += einsum("suv, sAtuv -> At", D_r_ao, F_1_ao_dfa)

        self.grad_gga = grad_contrib
        return self

    @timing
    def prepare_gradient_pt2(self):
        tensors = self.tensors
        C, D, e = self.C, self.D, self.e
        mol, aux_ri = self.mol, self.aux_ri
        natm, nao, nmo, nocc, nvir, naux = mol.natm, self.nao, self.nmo, self.nocc, self.nvir, self.df_ri.get_naoaux()
        mocc, mvir = max(nocc), max(nvir)
        # this algorithm asserts naux = aux.nao, i.e. no linear dependency in auxiliary basis
        assert naux == aux_ri.nao
        so, sv, sa = self.so, self.sv, self.sa

        D_r = tensors.load("D_r")
        H_1_mo = tensors.load("H_1_mo")
        grad_corr = einsum("spq, sApq -> A", D_r, H_1_mo)
        if not self.eval_pt2:
            grad_corr.shape = (natm, 3)
            self.grad_pt2 = grad_corr
            return

        W_I = tensors.load("W_I")
        W_II = - einsum("spq, sq -> spq", D_r, e)
        W_III_tmp = self.Ax0_Core(so, so, sa, sa)(D_r)
        W = W_I + W_II
        for σ in (α, β):
            W[σ][so[σ], so[σ]] += - 0.5 * W_III_tmp[σ]
        W_ao = einsum("sup, spq, svq -> suv", C, W, C)
        S_1_ao = tensors.load("S_1_ao")
        grad_corr += np.einsum("suv, Auv -> A", W_ao, S_1_ao)
        grad_corr.shape = (natm, 3)

        L_inv, L_1_gen = generator_L_1(aux_ri)
        int3c2e_ip1_gen = int3c_wrapper(mol, aux_ri, "int3c2e_ip1", "s1")
        int3c2e_ip2_gen = int3c_wrapper(mol, aux_ri, "int3c2e_ip2", "s1")
        Y_ia_ri = [np.asarray(tensors["Y_mo_ri" + str(σ)][:, so[σ], sv[σ]]) for σ in (α, β)]
        G_ia_ri = [tensors.load("G_ia_ri" + str(σ)) for σ in (α, β)]

        for A in range(natm):
            L_1_ri = L_1_gen(A)
            Y_1_ia_ri = [np.zeros((3, naux, nocc[σ], nvir[σ])) for σ in (α, β)]
            shA0, shA1, _, _ = mol.aoslice_by_atom()[A]
            shA0a, shA1a, _, _ = aux_ri.aoslice_by_atom()[A]

            nbatch = calc_batch_size(3*(nao+mocc)*naux, self.get_memory(), tot_size(Y_1_ia_ri))
            for shU0, shU1, U0, U1 in gen_shl_batch(mol, nbatch, shA0, shA1):
                su = slice(U0, U1)
                int3c2e_ip1 = int3c2e_ip1_gen((shU0, shU1, 0, mol.nbas, 0, aux_ri.nbas))
                for σ in (α, β):
                    Y_1_ia_ri[σ] -= einsum("tuvQ, PQ, ui, va -> tPia", int3c2e_ip1, L_inv, C[σ][su, so[σ]], C[σ][:, sv[σ]])
                    Y_1_ia_ri[σ] -= einsum("tuvQ, PQ, ua, vi -> tPia", int3c2e_ip1, L_inv, C[σ][su, sv[σ]], C[σ][:, so[σ]])

            nbatch = calc_batch_size(3*nao*(nao+mocc), self.get_memory(), tot_size(Y_1_ia_ri))
            for shP0, shP1, P0, P1 in gen_shl_batch(aux_ri, nbatch, shA0a, shA1a):
                sp = slice(P0, P1)
                int3c2e_ip2 = int3c2e_ip2_gen((0, mol.nbas, 0, mol.nbas, shP0, shP1))
                for σ in (α, β):
                    Y_1_ia_ri[σ] -= einsum("tuvQ, PQ, ui, va -> tPia", int3c2e_ip2, L_inv[:, sp], C[σ][:, so[σ]], C[σ][:, sv[σ]])

            for σ in (α, β):
                Y_1_ia_ri[σ] -= einsum("Qia, tRQ, PR -> tPia", Y_ia_ri[σ], L_1_ri, L_inv)
                grad_corr[A] += einsum("Pia, tPia -> t", G_ia_ri[σ], Y_1_ia_ri[σ])
        self.grad_pt2 = grad_corr

    @timing
    def prepare_gradient_enfunc(self):
        tensors = self.tensors
        natm = self.mol.natm
        Co, eo, D = self.Co, self.eo, self.D
        so = self.so

        grad_contrib = self.mf_s.Gradients().grad_nuc()
        grad_contrib.shape = (natm * 3,)

        H_1_ao = tensors.load("H_1_ao")
        S_1_mo = tensors.load("S_1_mo")

        grad_contrib += np.einsum("Auv, suv -> A", H_1_ao, D, optimize=True)  # TODO check PySCF lib.einsum why fails
        if self.xc_n is None:
            for σ in (α, β):
                grad_contrib -= np.einsum("Ai, i -> A", S_1_mo[σ][:, so[σ], so[σ]].diagonal(0, -1, -2), eo[σ])
        else:
            # TODO see whether get_fock could use mo_coeff to accelearate RI-K
            F_0_ao_n = self.mf_n.get_fock(dm=D)
            nc_F_0_ij = [(Co[σ].T @ F_0_ao_n[σ] @ Co[σ]) for σ in (α, β)]
            for σ in (α, β):
                grad_contrib -= einsum("Aij, ij -> A", S_1_mo[σ][:, so[σ], so[σ]], nc_F_0_ij[σ])
        grad_contrib.shape = (natm, 3)

        # handle dftd3 situation
        mol = self.mol
        if "D3" in self.xc_add:
            drv = itrf.libdftd3.wrapper_params
            params = np.asarray(self.xc_add["D3"][0], order="F")
            version = self.xc_add["D3"][1]
            coords = np.asarray(mol.atom_coords(), order="F")
            itype = np.asarray(mol.atom_charges(), order="F")
            edisp = np.zeros(1)
            grad = np.zeros((mol.natm, 3))
            drv(
                ctypes.c_int(mol.natm),  # natoms
                coords.ctypes.data_as(ctypes.c_void_p),  # coords
                itype.ctypes.data_as(ctypes.c_void_p),  # itype
                params.ctypes.data_as(ctypes.c_void_p),  # params
                ctypes.c_int(version),  # version
                edisp.ctypes.data_as(ctypes.c_void_p),  # edisp
                grad.ctypes.data_as(ctypes.c_void_p))  # grads)
            grad_contrib += grad

        self.grad_enfunc = grad_contrib

    def base_method(self) -> UDFDH:
        self.__class__ = UDFDH
        return self


# Licensed under a 3-clause BSD style license - see LICENSE.rst
# coding=utf-8

import subprocess
import warnings

import numpy as np
from skimage.transform import rotate as imrotate

import astropy.units as u
from astropy.io import ascii
from astropy.table import Table
from astropy import visualization

import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib.colors as col

from .config import recursive_subclasses, merge_config, mmtwfs_config
from .custom_exceptions import WFSConfigException
from .secondary import SecondaryFactory
from .zernike import ZernikeVector

import logging
import logging.handlers

# we need to wrap the poppy import in a context manager to trap its whinging about missing pysynphot stuff that we don't use.
with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    import poppy

log = logging.getLogger("Telescope")
log.setLevel(logging.INFO)

__all__ = ['TelescopeFactory', 'MMT', 'FLWO12']

def TelescopeFactory(telescope="mmt", secondary="f5", config={}, **kwargs):
    """
    Build and return proper Telescope sub-class instance based on the value of 'telescope' and 'secondary'.
    """
    config = merge_config(config, dict(**kwargs))
    secondary = secondary.lower()
    telescope = telescope.lower()

    types = recursive_subclasses(Telescope)
    telescopes = [t.__name__.lower() for t in types]
    tel_map = dict(list(zip(telescopes, types)))

    if telescope not in telescopes:
        raise WFSConfigException(value=f"Specified telescope, {telescope}, not valid or not implemented.")

    tel_cls = tel_map[telescope](secondary=secondary, config=config)
    return tel_cls

class Telescope(object):
    """
    Defines generic configuration and methods that pertain to telescope and primary mirror systems
    """
    def __init__(self, telescope="mmt", secondary="f5", config={}, **kwargs):
        config = merge_config(config, dict(**kwargs))
        if telescope not in mmtwfs_config['telescope']:
            msg = f"Invalid telescope specified, {telescope}."
            raise WFSConfigException(value=msg)
        if secondary not in mmtwfs_config['secondary']:
            msg = f"Invalid secondary specified, {secondary}."
            raise WFSConfigException(value=msg)
        if mmtwfs_config['secondary'][secondary]['telescope'] != telescope:
            msg = f"Invalid secondary, {secondary}, for telescope, {telescope}."
            raise WFSConfigException(value=msg)

        self.__dict__.update(merge_config(mmtwfs_config['telescope'][telescope], config))

        self.secondary = SecondaryFactory(secondary=secondary)

        self.radius = self.diameter / 2.
        self.nmperrad = self.radius.to(u.nm).value
        self.nmperasec = self.nmperrad / 206265.

        # ratio of the size of the central obstruction of the secondary to the size of the primary
        self.obscuration = self.secondary.diameter / self.diameter

        # create model of MMTO pupil including secondary and secondary support obstructions
        self.pupil = self._pupil_model()

        # initialize poppy optical system used for calculating the PSFs
        self.osys = poppy.OpticalSystem()
        self.osys.add_pupil(self.pupil)
        self.osys.add_pupil(poppy.ZernikeWFE(radius=self.radius.to(u.m).value, coefficients=[0.0, 0.0, 0.0, 0.0]))
        self.osys.add_detector(pixelscale=self.psf_pixel_scale, fov_arcsec=self.psf_fov)

    def _pupil_model(self):
        """
        Use poppy to create a model of the pupil given the configured primary and secondary mirrors.
        """
        primary = poppy.CircularAperture(radius=self.radius.to(u.m).value)
        secondary = poppy.SecondaryObscuration(
            secondary_radius=self.secondary.diameter.to(u.m).value / 2,
            n_supports=self.n_supports,
            support_width=self.support_width.to(u.m).value,
            support_angle_offset=self.support_offset.to(u.deg).value
        )
        pup_model = poppy.CompoundAnalyticOptic(opticslist=[primary, secondary], name="MMTO")
        return pup_model

    def pupil_mask(self, rotation=0.0, size=512):
        """
        Use the pupil model to make a pupil mask that can be used as a kernel for finding pupil-like things in images
        """
        if size >= 700:
            msg = "WFS pupil sizes are currently restricted to 700 pixels in diameter or less."
            raise WFSConfigException(value=msg)

        rotation = u.Quantity(rotation, u.deg)

        # not sure how to get the image data out directly, but the to_fits() method gives me a path...
        pup_im = imrotate(self.pupil.to_fits(npix=size)[0].data.astype(float), rotation.value)
        pup_im = pup_im / pup_im.max()
        return pup_im

    def psf(self, zv=ZernikeVector(), wavelength=550.*u.nm, plot=True):
        """
        Take a ZernikeVector and calculate resulting PSF at given wavelength.
        """
        # poppy wants the wavelength in meters
        try:
            w = wavelength.to(u.m).value
        except AttributeError:
            w = wavelength  # if no unit provided, assumed meters

        # poppy wants the piston term so whack it in there if modestart isn't already 1
        if zv.modestart != 1:
            zv.modestart = 1
            zv['Z01'] = 0.0

        # poppy wants coeffs in meters
        zv.units = u.m

        # poppy wants Noll normalized coefficients
        coeffs = zv.norm_array

        # pop detector out to reuse, pop old wavefront error out to make way for new
        det = self.osys.planes.pop()
        fov = det.fov_arcsec.value

        # add new wavefront error and put detector back in place
        wfe = poppy.ZernikeWFE(radius=self.radius.to(u.m).value, coefficients=coeffs)
        self.osys.add_pupil(wfe)
        self.osys.planes.append(det)

        psf = self.osys.calc_psf(w)

        psf_fig = None
        if plot:
            psf_fig, ax = plt.subplots()
            psf_fig.set_label("PSF at {0:0.0f}".format(wavelength))
            norm = visualization.mpl_normalize.ImageNormalize(stretch=visualization.LinearStretch())
            ims = ax.imshow(psf[0].data, extent=[-fov/2, fov/2, -fov/2, fov/2], cmap=cm.magma, norm=norm)
            ax.set_xlabel("arcsec")
            ax.set_ylabel("arcsec")
            cb = psf_fig.colorbar(ims)
            cb.set_label("Fraction of Total Flux")
        return psf, psf_fig

class FLWO12(Telescope):
    """
    Defines configuration and methods for the FLWO 1.2-meter
    """
    def __init__(self, config={}, **kwargs):
        config = merge_config(config, dict(**kwargs))
        super(FLWO12, self).__init__(telescope="flwo12", secondary="flwo12", config=config)

class FLWO15(Telescope):
    """
    Defines configuration and methods for the FLWO 1.5-meter
    """
    def __init__(self, config={}, **kwargs):
        config = merge_config(config, dict(**kwargs))
        super(FLWO15, self).__init__(telescope="flwo15", secondary="flwo15", config=config)

class MMT(Telescope):
    """
    Defines configuration and methods that pertain to the MMT's telescope and primary mirror systems
    """
    def __init__(self, secondary="f5", config={}, **kwargs):
        config = merge_config(config, dict(**kwargs))
        super(MMT, self).__init__(telescope="mmt", secondary=secondary, config=config)

        # load table of finite element coordinates
        self.nodecoor = self.load_bcv_coordinates()
        self.n_node = len(self.nodecoor)

        # load table of actuator coordinates
        self.actcoor = self.load_actuator_coordinates()
        self.n_act = len(self.actcoor)

        # load actuator influence matrix that provides the surface displacement caused by 1 N of force by
        # each actuator at each of self.node finite element node positions.
        self.surf2act = self.load_influence_matrix()

        # use this boolean to determine if corrections are actually to be sent
        self.connected = False

        # keep track of last and total forces. a blank ZernikeVector will generate the appropriate format
        # table with all forces set to 0.
        self.last_forces = self.bending_forces(zv=ZernikeVector())
        self.total_forces = self.bending_forces(zv=ZernikeVector())
        self.last_m1focus = 0.0 * u.um
        self.total_m1focus = 0.0 * u.um

    def connect(self):
        """
        Set state to connected so that calculated corrections will be sent to the relevant systems
        """
        self.connected = True

    def disconnect(self):
        """
        Set state to disconnected so that corrections will be calculated, but not sent
        """
        self.connected = False

    def bending_forces(self, zv=ZernikeVector(), gain=0.5):
        """
        Given a ZernikeVector (or similar object describing a 2D polynomial surface), calculate the actuator forces required
        to correct for the surface displacement it describes.
        """
        # we don't want to bend any tilts...
        if 'Z02' in zv:
            zv['Z02'] = 0.0
        if 'Z03' in zv:
            zv['Z03'] = 0.0

        # convert to nm...
        zv.units = u.nm

        # make sure we're not Noll normalized...
        zv.denormalize()

        # need to rotate the wavefront -90 degrees to match the BCV angle convention of +Y being 0 deg.
        zv.rotate(-90*u.deg)

        # get surface displacements at the BCV node positions. multiply the wavefront amplitude by 0.5 to account for
        # reflection off the surface.
        surf_corr = -0.5 * gain * zv.total_phase(self.nodecoor['bcv_rho'], self.nodecoor['bcv_phi'])
        if isinstance(surf_corr, float):  # means we got 0.0 from zv.total_phase()
            force_vec = np.zeros(self.n_act)
        else:
            force_vec = np.dot(surf_corr, self.surf2act).value  # remove the units that got passed through

        # return an astropy.table.Table so we can easily package actuator ID along with the force. its write() method
        # also provides a lot of flexibility in providing outputs that match the old system.
        t = Table([self.actcoor['act_id'], force_vec], names=['actuator', 'force'])
        return t

    def to_rcell(self, t, filename="zfile", overwrite=True):
        """
        Take table generated by bending_forces() and write it to a file of a format that matches the old SHWFS system
        """
        t.write(filename, format="ascii.no_header", delimiter="\t", formats={'force': ".1f"}, overwrite=overwrite)

    def calculate_primary_corrections(self, zv, mask=[], gain=0.5):
        """
        Take ZernikeVector as input and determine corrections to apply to primary/secondary
        """
        # leave out tilts, focus, and coma from force calcs to start with
        def_mask = ['Z02', 'Z03', 'Z04', 'Z07', 'Z08']
        def_mask.extend(mask)

        # mask out all high order terms beyond 2nd order spherical
        for i in range(23, 99):
            def_mask.append("Z{0:02d}".format(i))

        mask = list(set(def_mask))
        zv_masked = zv.copy()
        zv_masked.denormalize()
        for k in mask:
            zv_masked.ignore(k)

        # to reduce the amount of force required to remove spherical aberration, we offset the r**2 part of that term by
        # bending focus into the primary and then offsetting that by adjusting the secondary.  this has the effect of
        # reducing by ~1/4 to 1/3 the total force required to correct a given amount of spherical aberration.
        #
        # this same scheme can also be extended to the higher order spherical terms as well, Z22 and Z37.
        #
        # for reference:
        #   Z04 ~ 2r**2 - 1
        #   Z11 ~ 6r**4 - 6r**2 + 1
        #   Z22 ~ 20r**6 - 30r**4 + 12r**2 - 1
        #   Z37 ~ 70r**8 - 140r**6 + 90r**4 - 20r**2 + 1
        #
        zv_masked['Z04'] = -6.0 * zv_masked['Z11'] - 12.0 * zv_masked['Z22'] - 20.0 * zv_masked['Z37']

        m1focus_corr = gain * zv_masked['Z04'] / self.secondary.focus_trans

        t = self.bending_forces(zv=zv_masked, gain=gain)

        return t, m1focus_corr, zv_masked

    def bend_mirror(self, filename="zfile"):
        """
        Take a force file and send it to the cell to apply bending forces. Return fraction of requested
        forces that the cell reports were applied.
        """
        frac = 1.0
        log.info(f"Using command, /mmt/scripts/cell_send_forces {filename}, to apply forces...")
        pipe = subprocess.Popen(
            ['/mmt/scripts/cell_send_forces', f"{filename}"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )

        try:
            (stdout, stderr) = pipe.communicate(timeout=10)
        except subprocess.TimeoutExpired:
            pipe.kill()
            (stdout, stderr) = pipe.communicate()

        outstr = stdout.decode('utf8')
        outerr = stderr.decode('utf8')

        # had to dig into the cell code at /mmt/vxsource/mmt/cell/src/cell_inf.c to get the messages that are produces
        if "Able to Apply" in outstr:
            log.info("...forces successfully applied in full.")

        elif "Forces Rejected" in outstr:
            log.error("...forces rejected!")
            frac = 0.0

        elif "Unable to apply forces" in outstr:
            log.error("...unable to apply forces!")
            frac = 0.0

        elif "Applying partial forces" in outstr:
            percent = float(outstr.split()[3])
            frac = percent / 100.0
            log.warn(f"...applied {percent}% of requested forces.")

        else:
            log.error(f"...got unexpected reply from cell command: {outstr}")
            log.error(f"\t stderr: {outerr}")

        return frac

    def correct_primary(self, t, m1focus_corr, filename="zfile"):
        """
        Take force table and focus offset calculated by self.calculate_primary_corrections() and apply them, if connected.
        """
        frac = 1.0
        if self.connected:
            self.to_rcell(t, filename=filename)
            log.info(f"Sending forces from {filename}...")
            frac = self.bend_mirror(filename=filename)
            self.secondary.m1spherical(frac * m1focus_corr)
        else:
            log.info("Not connected; no commands sent to cell or hexapod.")

        self.last_forces = t.copy(copy_data=True)
        self.last_forces['force'] *= frac
        self.last_m1focus = frac * m1focus_corr.copy()
        self.total_forces['force'] += frac * t['force']
        self.total_m1focus += frac * m1focus_corr
        return t, m1focus_corr

    def undo_last(self, zfilename="zfile_undo"):
        """
        Undo the last set of corrections.
        """
        self.last_forces['force'] *= -1
        self.last_m1focus *= -1
        frac = 1.0
        if self.connected:
            log.info("Undoing last set of primary mirror corrections...")
            self.to_rcell(self.last_forces, filename=zfilename)
            frac = self.bend_mirror(filename=zfilename)
            self.secondary.m1spherical(frac * self.last_m1focus)
        else:
            log.info("Not connected; no undo commands sent.")

        self.total_m1focus += frac * self.last_m1focus
        self.total_forces['force'] += frac * self.last_forces['force']
        return self.last_forces.copy(), self.last_m1focus.copy()

    def clear_forces(self):
        """
        Clear applied forces from primary mirror and clear any m1spherical offsets from secondary hexapod
        """
        if self.connected:
            log.info("Clearing forces and spherical aberration focus offsets...")
            self.secondary.clear_m1spherical()
            pipe = subprocess.Popen(['/mmt/scripts/cell_clear_forces'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)

            try:
                (stdout, stderr) = pipe.communicate(timeout=20)
            except subprocess.TimeoutExpired:
                pipe.kill()
                (stdout, stderr) = pipe.communicate()

            outstr = stdout.decode('utf8')
            outerr = stderr.decode('utf8')
            log.info(f"...{outstr.strip()}")
            if len(outerr) > 0:
                log.warn(f"Got error from cell_clear_forces: {outerr}")
        else:
            log.info("Not connected; no clearing commands sent.")

        # the 'last' corrections are negations of the current total. reset the totals to 0.
        self.last_forces = self.total_forces.copy(copy_data=True)
        self.last_forces['force'] *= -1
        self.last_m1focus = -self.total_m1focus
        self.total_forces = self.bending_forces(zv=ZernikeVector())
        self.total_m1focus = 0.0

        return self.last_forces, self.last_m1focus

    def load_influence_matrix(self):
        """
        The influence of each actuator on the mirror surface has been modeled via finite element analysis.
        This method loads the influence matrix that resulted from this analysis and maps for each actuator
        the influence of 1 lb of force on the mirror surface at each finite element node.  This matrix is
        stored in a binary file for compactness and speed of loading.
        """
        surf2act = np.fromfile(self.surf2act_file, dtype=np.float32).reshape(self.n_act, self.n_node).transpose()
        return surf2act

    def load_actuator_coordinates(self):
        """
        The actuator IDs and X/Y positions in mm are stored in a simple ASCII table.  Load it using
        astropy.io.ascii, convert to units of mirror radius, and add polar coordinates.
        """
        coord = ascii.read(self.actuator_file, names=["act_id", "act_x", "act_y", "act_type"])
        for ax in ["act_x", "act_y"]:
            coord[ax] /= self.bcv_radius.to(u.mm).value
        coord['act_rho'] = np.sqrt(coord['act_x']**2 + coord['act_y']**2)
        coord['act_phi'] = np.arctan2(coord['act_y'], coord['act_x'])
        coord['act_phi'].unit = u.radian

        return coord

    def load_bcv_coordinates(self):
        """
        The BCV finite element nodes IDs and X/Y/Z positions in mm are stored in a simple ASCII table.  Load it
        using astropy.io.ascii, convert to units of mirror radius, and add polar coordinates.
        """
        coord = ascii.read(self.nodecoor_file, names=["bcv_id", "bcv_x", "bcv_y", "bcv_z"])
        for ax in ["bcv_x", "bcv_y"]:
            coord[ax] /= self.bcv_radius.to(u.mm).value
        coord['bcv_rho'] = np.sqrt(coord['bcv_x']**2 + coord['bcv_y']**2)
        coord['bcv_phi'] = np.arctan2(coord['bcv_y'], coord['bcv_x'])
        coord['bcv_phi'].unit = u.radian

        return coord

    def plot_forces(self, t, m1focus=None, limit=100.):
        """
        Plot actuator forces given force table as output from self.bending_forces()
        """
        coords = self.actcoor
        r_fac = 0.5 * self.diameter / self.bcv_radius  # adjust for slight difference
        cmap = cm.ScalarMappable(col.Normalize(-1*limit, limit), cm.bwr)
        cmap._A = []  # grr stupid matplotlib
        fig, ax = plt.subplots()
        fig.set_label("M1 Actuator Forces")
        xcor, ycor = coords['act_x']/r_fac, coords['act_y']/r_fac
        ax.scatter(xcor, ycor, color=cmap.to_rgba(t['force']))
        for i, (x, y) in enumerate(zip(xcor, ycor)):
            ax.text(x, y+0.02, t['actuator'][i],  horizontalalignment='center', verticalalignment='bottom', size='xx-small')

        ax.set_aspect(1.0)
        circle1 = plt.Circle((0, 0), 1.0, fill=False, color='black', alpha=0.2)
        circle2 = plt.Circle((0, 0), 0.9/6.5, fill=False, color='black', alpha=0.2)
        ax.add_artist(circle1)
        ax.add_artist(circle2)
        if m1focus is not None:
            ax.set_title("M1 Focus Offset: {0:0.1f}".format(m1focus))
        ax.set_axis_off()
        cb = fig.colorbar(cmap)
        cb.set_label("Actuator Force (N)")
        return fig

import numpy,import astropy,from astropy


#!/usr/bin/env python
import os
import numpy as np
import time
import copy
import sys

import argparse

ang_2_bohr = 1.0/0.52917721067
hart_2_ev = 27.21138602

import cp2k_spm_tools.cp2k_grid_orbitals as cgo
import cp2k_spm_tools.cp2k_stm_sts as css
from cp2k_spm_tools import common, cube_utils
from cp2k_spm_tools.cube import Cube

from mpi4py import MPI

comm = MPI.COMM_WORLD
mpi_rank = comm.Get_rank()
mpi_size = comm.Get_size()

parser = argparse.ArgumentParser(
    description='Puts the CP2K orbitals on grid and calculates STM.')

### ----------------------------------------------------------------------
### Input and output files
parser.add_argument(
    '--cp2k_input_file',
    metavar='FILENAME',
    required=True,
    help='CP2K input of the SCF calculation.')
parser.add_argument(
    '--basis_set_file',
    metavar='FILENAME',
    required=True,
    help='File containing the used basis sets.')
parser.add_argument(
    '--xyz_file',
    metavar='FILENAME',
    required=True,
    help='.xyz file containing the geometry.')
parser.add_argument(
    '--wfn_file',
    metavar='FILENAME',
    required=True,
    help='Restart file containing the final wavefunction.')
parser.add_argument(
    '--hartree_file',
    metavar='FILENAME',
    required=True,
    help='Cube file containing the hartree potential.')
parser.add_argument(
    '--output_file',
    metavar='FILENAME',
    default="./stm.npz",
    help='File, where to save the STM/STS output')
parser.add_argument(
    '--orb_output_file',
    metavar='FILENAME',
    default="./orb.npz",
    help='File, where to save the orbital output')
### ----------------------------------------------------------------------
### Parameters for putting orbitals on grid
parser.add_argument(
    '--eval_region',
    type=str,
    nargs=6,
    metavar='X',
    required=True,
    help=common.eval_region_description
)
parser.add_argument(
    '--dx',
    type=float,
    metavar='DX',
    required=True,
    help='Spatial step for the grid (angstroms).')
parser.add_argument(
    '--eval_cutoff',
    type=float,
    metavar='D',
    default=16.0,
    help=("Size of the region around the atom where each"
          " orbital is evaluated (only used for 'G' region).")
)
parser.add_argument(
    '--extrap_extent',
    type=float,
    metavar='H',
    default=4.0,
    required=True,
    help="The extent of the extrapolation region. (angstrom)"
)
### ----------------------------------------------------------------------
### Gas phase analysis parameters - image at orbital energies
parser.add_argument(
    '--n_homo',
    type=int,
    metavar='N',
    default=0,
    help="Number of HOMO orbitals to analyse.")
parser.add_argument(
    '--n_lumo',
    type=int,
    metavar='N',
    default=0,
    help="Number of LUMO orbitals to analyse.")
parser.add_argument(
    '--orb_heights',
    nargs='*',
    type=float,
    metavar='H',
    help="List of heights for constant height orbital pictures (wrt topmost atom).")
parser.add_argument(
    '--orb_isovalues',
    nargs='*',
    type=float,
    metavar='C',
    help="List of charge density isovalues for constant current orbital pictures")
parser.add_argument(
    '--orb_fwhms',
    nargs='*',
    type=float,
    default=[0.02],
    help="Full width at half maximum for orbital STS gaussian broadening. (eV)")
### ----------------------------------------------------------------------
### Slab system analysis parameters - images at specified energies
###
### Option 1: continuous selection
parser.add_argument(
    '--energy_range',
    nargs=3,
    type=float,
    metavar='E',
    help='Selection of STM/STS energy values based on a range: min, max and differential.')
###
### Option 2: discrete selection
parser.add_argument(
    '--energies',
    nargs='*',
    type=float,
    metavar='E',
    help='Discrete energies where to run the STM/STS.')
### ----------------------------------------------------------------------
### Parameters for STM/STS series
parser.add_argument(
    '--heights',
    nargs='*',
    type=float,
    metavar='H',
    help="List of heights for constant height STM pictures (wrt topmost atom).")
parser.add_argument(
    '--isovalues',
    nargs='*',
    type=float,
    metavar='C',
    help="List of charge density isovalues for constant current STM pictures.")
parser.add_argument(
    '--fwhms',
    nargs='*',
    type=float,
    default=[0.1],
    help="Full width at half maximum for STS gaussian broadening. (eV)")
### ----------------------------------------------------------------------
### P - tip ratio list
parser.add_argument(
    '--p_tip_ratios',
    nargs='+',
    type=float,
    metavar='P',
    default=[0.0],
    help=("List of p character of the STM tip: 0.0 corresponds"
          "to fully s-type and 1.0 to fully p-type tip")
)
### ----------------------------------------------------------------------

time0 = time.time()

### ------------------------------------------------------
### Parse args for only one rank to suppress duplicate stdio
### ------------------------------------------------------

args = None
args_success = False
try:
    if mpi_rank == 0:
        args = parser.parse_args()
        args_success = True
finally:
    args_success = comm.bcast(args_success, root=0)

if not args_success:
    print(mpi_rank, "exiting")
    exit(0)

args = comm.bcast(args, root=0)

### ------------------------------------------------------
### Energy values for STM/STS
### ------------------------------------------------------

if args.energies is not None:
    e_arr = np.array(args.energies)
elif args.energy_range is not None:
    emin, emax, de = args.energy_range
    e_arr = np.arange(emin, emax+de/2, de)
else:
    e_arr = None

max_fwhm = np.max(args.fwhms)
if e_arr is not None:
    sel_emin = np.min(e_arr) - 2.0*max_fwhm
    sel_emax = np.max(e_arr) + 2.0*max_fwhm
else:
    sel_emin = None
    sel_emax = None

### ------------------------------------------------------
### Evaluate orbitals on the real-space grid
### ------------------------------------------------------

cp2k_grid_orb = cgo.Cp2kGridOrbitals(mpi_rank, mpi_size, mpi_comm=comm, single_precision=True)
cp2k_grid_orb.read_cp2k_input(args.cp2k_input_file)
cp2k_grid_orb.read_xyz(args.xyz_file)
cp2k_grid_orb.center_atoms_to_cell()
cp2k_grid_orb.read_basis_functions(args.basis_set_file)
cp2k_grid_orb.load_restart_wfn_file(args.wfn_file,
                                    emin=sel_emin, emax=sel_emax,
                                    n_occ=args.n_homo, n_virt=args.n_lumo
)

print("R%d/%d: loaded wfn, %.2fs"%(mpi_rank, mpi_size, (time.time() - time0)))
sys.stdout.flush()
time1 = time.time()

eval_reg = common.parse_eval_region_input(args.eval_region, cp2k_grid_orb.ase_atoms, cp2k_grid_orb.cell)

# --------
# Make sure extrap extent is compatible with heights
atoms_max_z = np.max(cp2k_grid_orb.ase_atoms.positions[:, 2])
eval_z_above_atoms = eval_reg[2][1] - atoms_max_z
extrap_extent = args.extrap_extent
for hs in [args.orb_heights, args.heights]:
    if hs is not None:
        if np.max(hs) - eval_z_above_atoms > extrap_extent:
            print("Increasing extrap. extent to be compatible with heights.")
            extrap_extent = np.max(hs)- eval_z_above_atoms
# --------

cp2k_grid_orb.calc_morbs_in_region(args.dx,
                                x_eval_region = eval_reg[0],
                                y_eval_region = eval_reg[1],
                                z_eval_region = eval_reg[2],
                                pbc = (True, True, False),
                                reserve_extrap = extrap_extent,
                                eval_cutoff = args.eval_cutoff)

print("R%d/%d: evaluated wfn, %.2fs"%(mpi_rank, mpi_size, (time.time() - time1)))
sys.stdout.flush()
time1 = time.time()

### ------------------------------------------------------
### Extrapolate orbitals
### ------------------------------------------------------

hart_cube = Cube()
hart_cube.read_cube_file(args.hartree_file)
extrap_plane_z = eval_reg[2][1] / ang_2_bohr - np.max(cp2k_grid_orb.ase_atoms.positions[:, 2])
hart_plane = hart_cube.get_plane_above_topmost_atom(extrap_plane_z) - cp2k_grid_orb.ref_energy/hart_2_ev

cp2k_grid_orb.extrapolate_morbs(hart_plane=hart_plane)

print("R%d/%d: extrapolated wfn, %.2fs"%(mpi_rank, mpi_size, (time.time() - time1)))
sys.stdout.flush()
time1 = time.time()

### ------------------------------------------------------
### Calculate the ionization potential (just for output)
### ------------------------------------------------------

if mpi_rank == 0:
    # NB: currently only accurate for isolated molecules
    if cp2k_grid_orb.nspin == 1:
        homo_en = cp2k_grid_orb.global_morb_energies[0][cp2k_grid_orb.i_homo_glob[0]]
    else:
        homo_en = np.max([
            cp2k_grid_orb.global_morb_energies[0][cp2k_grid_orb.i_homo_glob[0]],
            cp2k_grid_orb.global_morb_energies[1][cp2k_grid_orb.i_homo_glob[1]]
        ])
    ion_pot = cube_utils.find_vacuum_level_naive(hart_cube) - (homo_en + cp2k_grid_orb.ref_energy)
    print("IONIZATION POTENIAL (eV): %.6f (accurate only for isolated molecules)" % ion_pot)

### ------------------------------------------------------
### Set up STM object
### ------------------------------------------------------

stm = css.STM(mpi_comm = comm, cp2k_grid_orb = cp2k_grid_orb, p_tip_ratios = args.p_tip_ratios)
stm.gather_global_energies()
stm.divide_by_space()

### ------------------------------------------------------
### Run STM-STS analysis for orbitals
### ------------------------------------------------------

orb_heights = args.orb_heights if args.orb_heights is not None else []
orb_isovalues = args.orb_isovalues if args.orb_isovalues is not None else []
orb_fwhms = args.orb_fwhms if args.orb_fwhms is not None else []

if len(orb_fwhms) != 0 and (len(orb_heights) != 0 or len(orb_isovalues) != 0):

    orbital_list = list(range(-args.n_homo + 1, args.n_lumo + 1))

    stm.create_orb_series(orbital_list, orb_heights, orb_isovalues, orb_fwhms)

    stm.collect_and_save_orb_maps(path=args.orb_output_file)

### ------------------------------------------------------
### Run STM-STS analysis for general energies
### ------------------------------------------------------

heights = args.heights if args.heights is not None else []
isovalues = args.isovalues if args.isovalues is not None else []
fwhms = args.fwhms if args.fwhms is not None else []

if e_arr is not None and len(fwhms) != 0 and (len(heights) != 0 or len(isovalues) != 0):

    stm.calculate_stm_maps(fwhms, isovalues, heights, e_arr)

    stm.collect_and_save_stm_maps(path=args.output_file)

print("R%d/%d: finished, total time: %.2fs"%(mpi_rank, mpi_size, (time.time() - time0)))