#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
-------------------------------------------------------------------------------
Python Script for Training the Deep Learning Model (StressNet) for
Transpiration Stress (Short Vegetation)
-------------------------------------------------------------------------------
Author: Akash Koppa
Affiliation: Hydro-Climate Extremes Lab (H-CEL), Ghent University, Belgium
Contact:  akash.koppa@ugent.be
-------------------------------------------------------------------------------
Reference:
Koppa, A., Rains, D., Hulsman, P., Poyatos, R., and Miralles, D. G.
A deep learning-based hybrid model of global terrestrial evaporation.
Nature Communications 13, 1912 (2022). https://doi.org/10.1038/s41467-022-29543-7
-------------------------------------------------------------------------------
"""

## import libraries
import tensorflow as tf
import os as os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

## user defined configuration
inpdir = "<< Specify path to input data here >>"

# station
flxnet = {"stn": (os.path.join(inpdir, "sites_short_vegetation.h5"))} # sites
# input features (absolute values)
absfil = {"ate": (os.path.join(inpdir, "ate_short_vegetation.h5")), # air temperature
          "co2": (os.path.join(inpdir, "co2_short_vegetation.h5")), # carbon di oxide
          "ssh": (os.path.join(inpdir, "ssh_short_vegetation.h5")), # plant available water
          "swi": (os.path.join(inpdir, "swi_short_vegetation.h5")), # incoming shortwave radiation
          "vod": (os.path.join(inpdir, "vod_short_vegetation.h5")), # vegetation optical depth
          "vpd": (os.path.join(inpdir, "vpd_short_vegetation.h5"))} # vapor pressure deficit

# input features (anomaly values)
anmfil = {"ate": (os.path.join(inpdir, "ate_short_vegetation_anomaly.h5")),
          "co2": (os.path.join(inpdir, "co2_short_vegetation_anomaly.h5")),
          "ssh": (os.path.join(inpdir, "ssh_short_vegetation_anomaly.h5")),
          "swi": (os.path.join(inpdir, "swi_short_vegetation_anomaly.h5")),
          "vod": (os.path.join(inpdir, "vod_short_vegetation_anomaly.h5")),
          "vpd": (os.path.join(inpdir, "vpd_short_vegetation_anomaly.h5"))}

# target variable
tarfil = {"str": (os.path.join(inpdir, "str_short_vegetation.h5"))} # transpiration stress

# output path for the final trained StressNet
outdir = "<< Specify path to input data here >>"
outfil = os.path.join(outdir, "stressnet_short_vegetation")

## main code
def main():
    """
    Main control script

    Returns
    -------
    A trained machine learning model with the required stress formulation
    """

    ## get the list of stations to subset the input data
    # read in the fluxnet site locations
    flxsit = pd.read_hdf(path_or_buf = flxnet["stn"],   key="siteda")
    flxsit = flxsit.dropna(how = "any")
    sitreq = flxsit.index

    ## create a combined tensorflow dataset
    trndat, tstdat, sclmin, sclmax = h5totf(filabs = absfil,
                                            filanm = anmfil,
                                            filtar = tarfil,
                                            sitreq = sitreq,
                                            shufle = True,
                                            batchn = 100,
                                            trnper = 85)

    ## get the required model
    tstmod = funmod(inpshp = 12,
                    losobj = kge,
                    metric = kge,
                    optmod = tf.keras.optimizers.Adam(learning_rate = 0.000142))

    ## train the deep learning model
    histst = tstmod.fit(trndat,
                        epochs = 900,
                        validation_data = tstdat)

    ## plot evolution of the loss function
    evolut = pd.DataFrame(histst.history)
    evorms = evolut[["kge","val_kge"]]
    evorms.plot()

    ## save the model
    tstmod.save(outfil)

## functions
# function to preprocess input data
def h5totf(filabs, filanm, filtar, sitreq, shufle, batchn, trnper):
    """
    Script to convert the input hdf5 files into tensorflow datasets which act
    as the primary input for the deep neural network

    Parameters
    ----------
    filabs : dictionary
        full paths to input feature data (absolute values)
    filanm : dictionary
        full path to input feature data (anomalies)
    filtar : dictionary
        full path to the target variable
    sitreq : list
        list of fluxnet sites
    shufle : Boolean
        True if data needs to be shuffled
    batchn : integer
        number of batches into which the data needs to divided into
    trnper : integer
        percentage between 0 and 100 based on which the data will be
        divided into training and testing dataset

    Returns
    -------
    retn01 : A tf.dataset with the training input features and target variables
    retn02 : A tf.dataset with the testing input features and target variables
    retn03 : A vector with min values of the input features and target variables
    retn04 : A vector with max values of the input features and target variables

    """

    # loop through the different input variables and create final dictionary
    inpdat = {}
    for i in filabs.keys():
        print("var under process: " + i)
        # read in the absolute values
        tmpabs = pd.read_hdf(filabs[i])
        #  read in the anomaly values
        tmpanm = pd.read_hdf(filanm[i])

        # store the data in dictionary
        inpdat[i] = tmpabs
        inpdat[i+"anm"] = tmpanm
    # target variable
    tardat = {}
    for i in filtar.keys():
        print("var under process: " + i)
        tmptar = pd.read_hdf(filtar[i])
        tmptar[tmptar > 1] = 1.0
        tardat[i] = tmptar

    print(inpdat.keys())

    # loop through the stations and create a final dataset containing only
    #   the input variables
    datlst = []
    for i in sitreq:
        print("site under process: " + i)
        lsttmp = []
        for j in inpdat.keys():
            tmpvar = inpdat[j][i]
            tmpvar.name = j
            lsttmp.append(inpdat[j][i])
        # create a data frame of variables for each
        # append the target stress
        tartmp = tardat[list(tardat.keys())[0]][i]
        tartmp.name = list(tardat.keys())[0]
        lsttmp.append(tartmp)
        dattmp = pd.concat(lsttmp, axis = 1)
        dattmp = dattmp.replace(to_replace = -9999.0, value = np.nan)
        dattmp = dattmp.replace(to_replace = -999.0, value = np.nan )
        dattmp = dattmp.dropna(axis = 0, how = "any")
        datlst.append(dattmp)

    del dattmp

    # create the final pandas data frame
    datfin = pd.concat(datlst, ignore_index = True)
    datfin = datfin.dropna(axis = 0, how = "any")
    print(datfin.shape)
    print(datfin.columns)

    # reset index
    datfin = datfin.reset_index(drop = True)

    # shuffle the data
    if shufle == True:
        print(" >> shuffling data")
        datfin = datfin.sample(frac = 1)

    # normalize by max
    print(" >> normalizing the data by max")
    tarval = datfin.pop(list(tardat.keys())[0])

    # calculate the max and min for backup
    datmin = datfin.quantile(0.05)
    datmax = datfin.quantile(0.95)
    datfin = datfin/datmax

    # convert the pandas data frame to a tf.dataset
    print(" >> Converting data frame into a tensorflow dataset")
    print(datfin.columns)
    print(datfin.shape)
    fullda = tf.data.Dataset.from_tensor_slices((datfin.values,
                                                 tarval.values))

    # split into training and testing datasets
    trnsiz = int((trnper/100) * len(datfin))
    # training dataset
    retn01 = fullda.take(trnsiz)
    # testing dataset
    retn02 = fullda.skip(trnsiz)

    if shufle == True:
        retn01 = retn01.shuffle(trnsiz)
        tstsiz = len(datfin) - trnsiz
        retn02 = retn02.shuffle(tstsiz)

    retn01 = retn01.batch(batchn)
    retn02 = retn02.batch(batchn)
    retn03 = datmin
    retn04 = datmax

    return retn01, retn02, retn03, retn04

# kling gupta efficiency
def kge(actual, predct):
    """
    a custom loss function based on the Kling Gupta Efficiency
    formula: [1 - sqrt((r-1)**2 + ((stddev_sim/stddev_obs)-1)**2 + ((mean_sim/mean_obs) - 1)**2)]
    reference: Decomposition of the mean squared error and NSE performance criteria:
               Implications for improving hydrological modelling. (2009).
               Journal of Hydrology, 377(1–2), 80–91.
               DOI: https://doi.org/10.1016/j.jhydrol.2009.08.003

    Parameters
    ----------
    actual : tensor
        ground truth data for the predictions to be compared against
    predct : tensor
        predicted data

    Returns
    -------
    (1-kge): scalar
        The loss function to be minimized

    """
    # >>> correlation
    acmean = tf.math.reduce_mean(actual)
    pdmean = tf.math.reduce_mean(predct)
    acmdev, pdmdev = actual - acmean, predct - pdmean
    cornum = tf.math.reduce_mean(tf.multiply(acmdev, pdmdev))
    corden = tf.math.reduce_std(acmdev) * tf.math.reduce_std(pdmdev)
    corcof = cornum / corden
    cratio = (corcof - 1)**2

    # variability ratio
    actstd = tf.math.reduce_std(actual)
    prestd = tf.math.reduce_std(predct)
    stdrat = prestd / actstd
    vratio = (stdrat - 1)**2

    # bias ratio (Beta)
    menrat = pdmean / acmean
    bratio = (menrat - 1)**2

    kgeval = 1 - tf.math.sqrt(cratio + vratio + bratio)
    retn01 = 1 - kgeval

    return retn01

# deep learning model
def funmod(inpshp, losobj, metric, optmod):
    """
    Function to create a machine learning model using the Functional API module
    of tensorflow. This module can be used to create more powerful machine
    learning models compared to the sequential module

    Parameters
    ----------
    inpshp : integer
        The number of input variables
    losobj : tf.keras.losses
        A tensorflow loss function
    metric : tf.keras.metrics
        A tensorflow error metric function
    optmod : character
        An optimizer which is available in tensorflow

    Returns
    -------
    retn01 : tf.keras.Model
        A compiled Functional API model

    """
    print(" >>> creating a Functional API model")
    # define the input layers
    inplyr = tf.keras.Input(shape = (inpshp, ))

    crslyr = tf.keras.layers.Dense(512, activation=tf.nn.swish)(inplyr)
    crslyr = tf.keras.layers.Dropout(0.45)(crslyr)
    crslyr = tf.keras.layers.Dense(256, activation=tf.nn.swish)(crslyr)
    crslyr = tf.keras.layers.Dropout(0.3)(crslyr)

    seqlyr = tf.keras.layers.Dense(792, activation=tf.nn.swish)(inplyr)
    seqlyr = tf.keras.layers.Dropout(0.45)(seqlyr)
    seqlyr = tf.keras.layers.Dense(512, activation=tf.nn.gelu)(seqlyr)
    seqlyr = tf.keras.layers.Dropout(0.45)(seqlyr)

    concat = tf.keras.layers.concatenate([crslyr, seqlyr, inplyr])

    outlyr = tf.keras.layers.Dense(768, activation=tf.nn.swish)(concat)
    outlyr = tf.keras.layers.Dropout(0.4)(outlyr)
    outlyr = tf.keras.layers.Dense(384, activation=tf.nn.swish)(outlyr)
    outlyr = tf.keras.layers.Dropout(0.4)(outlyr)
    outlyr = tf.keras.layers.Dense(256, activation=tf.nn.swish)(outlyr)
    outlyr = tf.keras.layers.Dropout(0.4)(outlyr)
    conca1 = tf.keras.layers.concatenate([outlyr, inplyr])
    outly1 = tf.keras.layers.Dense(256, activation=tf.nn.swish)(conca1)
    outly1 = tf.keras.layers.Dropout(0.3)(outly1)
    outly1 = tf.keras.layers.Dense(172, activation=tf.nn.swish)(outlyr)
    outly1 = tf.keras.layers.Dropout(0.35)(outlyr)
    conca2 = tf.keras.layers.concatenate([outly1, inplyr])

    outly2 = tf.keras.layers.Dense(128, activation=tf.nn.swish)(conca2)
    outly2 = tf.keras.layers.Dropout(0.3)(outly2)
    outly2 = tf.keras.layers.Dense(64, activation=tf.nn.gelu)(outly2)
    outmod = tf.keras.layers.Dense(6)(outly2)

    # combine the layers into a full model
    tmpmod = tf.keras.Model(inputs = inplyr,
                            outputs = outmod,
                            name = "glmfun")

    # define a loss function
    losobj = losobj
    # define a metric function
    metric = [metric]
    # compile the model
    tmpmod.compile(optimizer = optmod,
                   loss = losobj,
                   metrics = metric)
    # return the compiled model
    retn01 = tmpmod

    return retn01

## run the main script
if __name__ == "__main__":
    main()


from __future__ import annotations
import astropy.coordinates
import astropy.units
import astropy.time
import collections
import datetime
import dataclasses
import erfa
import functools
import itertools
import logging
import numpy
import pathlib
import requests
import scipy.optimize
import typing
from . import provider as provider
from . import satellite as satellite
from . import timesystem as timesystem
from .parse import (
    Version as Version,
    FileType as FileType,
    Record as Record,
    Satellite as Satellite,
    Product as Product,
)
from .provider import cddis as cddis

@dataclasses.dataclass(init=False)
class Id:
    value: bytes

    def __init__(self, _: typing.Union[str, bytes]):
        raise NotImplementedError()

@dataclasses.dataclass(init=False)
class Sp3Id(Id):
    def __init__(self, value: typing.Union[str, bytes]):
        self.value = value.encode() if isinstance(value, str) else value
        assert satellite.sp3_pattern.match(self.value) is not None

@dataclasses.dataclass(init=False)
class NoradId(Id):
    def __init__(self, value: typing.Union[str, bytes, int]):
        if isinstance(value, int):
            self.value = str(value).encode()
        elif isinstance(value, str):
            self.value = value.encode()
        else:
            self.value = value
        assert satellite.norad_pattern.match(self.value) is not None

total_seconds = numpy.vectorize(lambda delta: delta.total_seconds())

@dataclasses.dataclass
class PiecewisePolynomial:
    minimum_time: datetime.datetime
    maximum_time: datetime.datetime
    reference_time: datetime.datetime
    offset: numpy.ndarray
    begin: numpy.ndarray
    coefficients: numpy.ndarray
    velocity_coefficients: numpy.ndarray

    def __call__(self, obstime: astropy.time.Time) -> astropy.coordinates.ITRS:
        assert obstime.scale == "utc"
        first_obstime: datetime.datetime = obstime.min().to_datetime(timezone=datetime.timezone.utc)  # type: ignore
        if first_obstime < self.minimum_time:
            raise Exception(
                "the first obstime is too close to the first record to interpolate"
            )
        last_obstime: datetime.datetime = obstime.max().to_datetime(timezone=datetime.timezone.utc)  # type: ignore
        if last_obstime >= self.maximum_time:
            raise Exception(
                "the last obstime is too close to the last record to interpolate"
            )
        relative_obstime = total_seconds(
            obstime.to_datetime(timezone=datetime.timezone.utc) - self.reference_time
        )
        index = numpy.searchsorted(self.begin, relative_obstime, side="right") - 1
        relative_obstime -= self.offset[index]
        return astropy.coordinates.ITRS(
            x=(
                (
                    numpy.polynomial.polynomial.polyval(
                        relative_obstime,
                        self.coefficients[:, index, 0],
                        tensor=False,
                    )
                )
            )
            * astropy.units.Unit("m"),
            y=(
                (
                    numpy.polynomial.polynomial.polyval(
                        relative_obstime,
                        self.coefficients[:, index, 1],
                        tensor=False,
                    )
                )
            )
            * astropy.units.Unit("m"),
            z=(
                (
                    numpy.polynomial.polynomial.polyval(
                        relative_obstime,
                        self.coefficients[:, index, 2],
                        tensor=False,
                    )
                )
            )
            * astropy.units.Unit("m"),
            v_x=(
                (
                    numpy.polynomial.polynomial.polyval(
                        relative_obstime,
                        self.velocity_coefficients[:, index, 0],
                        tensor=False,
                    )
                )
            )
            * astropy.units.Unit("m/s"),
            v_y=(
                (
                    numpy.polynomial.polynomial.polyval(
                        relative_obstime,
                        self.velocity_coefficients[:, index, 1],
                        tensor=False,
                    )
                )
            )
            * astropy.units.Unit("m/s"),
            v_z=(
                (
                    numpy.polynomial.polynomial.polyval(
                        relative_obstime,
                        self.velocity_coefficients[:, index, 2],
                        tensor=False,
                    )
                )
            )
            * astropy.units.Unit("m/s"),
            obstime=obstime,
        )

def narrowed_records_to_piecewise_polynomial(
    records: typing.Sequence[Record],
    window: int,
    degree: int,
) -> PiecewisePolynomial:
    assert degree >= 0
    assert window * 2 + 1 > degree
    if len(records) < window * 2 + 1:
        raise Exception("insufficient number of records")
    record_position = numpy.array(tuple(record.position for record in records))
    relative_record_time = total_seconds(
        numpy.array(tuple(record.time - records[0].time for record in records))
    )
    relative_record_time_diff = numpy.diff(relative_record_time)
    offset = relative_record_time[window : len(records) - window]
    begin = (
        relative_record_time[window : len(records) - window]
        - relative_record_time_diff[window - 1 : len(records) - window - 1] / 2
    )
    coefficients = numpy.zeros(
        (len(records) - 2 * window, degree + 1, 3), dtype=numpy.double
    )
    mean = numpy.zeros((len(records) - 2 * window, 3))
    std = numpy.zeros((len(records) - 2 * window, 3))
    for index in range(0, len(coefficients)):
        mean[index] = numpy.mean(
            record_position[index : index + 2 * window + 1], axis=0
        )
        std[index] = numpy.std(record_position[index : index + 2 * window + 1], axis=0)
        coefficients[index] = numpy.polynomial.polynomial.polyfit(
            relative_record_time[index : index + 2 * window + 1] - offset[index],
            (record_position[index : index + 2 * window + 1] - mean[index])
            / std[index],
            deg=degree,
        )
    coefficients = coefficients.transpose((1, 0, 2)) * std
    coefficients[0] += mean
    if all(record.velocity is not None for record in records):
        record_velocity = numpy.array(tuple(record.velocity for record in records))
        velocity_coefficients = numpy.zeros(
            (len(records) - 2 * window, degree + 1, 3), dtype=numpy.double
        )
        velocity_mean = numpy.zeros((len(records) - 2 * window, 3))
        velocity_std = numpy.zeros((len(records) - 2 * window, 3))
        for index in range(0, len(velocity_coefficients)):
            velocity_mean[index] = numpy.mean(
                record_velocity[index : index + 2 * window + 1], axis=0
            )
            velocity_std[index] = numpy.std(
                record_velocity[index : index + 2 * window + 1], axis=0
            )
            velocity_coefficients[index] = numpy.polynomial.polynomial.polyfit(
                relative_record_time[index : index + 2 * window + 1] - offset[index],
                (record_velocity[index : index + 2 * window + 1] - velocity_mean[index])
                / velocity_std[index],
                deg=degree,
            )
        velocity_coefficients = (
            velocity_coefficients.transpose((1, 0, 2)) * velocity_std
        )
        velocity_coefficients[0] += velocity_mean
    else:
        velocity_coefficients = numpy.polynomial.polynomial.polyder(coefficients)
    return PiecewisePolynomial(
        minimum_time=records[window].time,
        maximum_time=records[-window].time,
        reference_time=records[0].time,
        offset=offset,
        begin=begin,
        coefficients=coefficients,
        velocity_coefficients=velocity_coefficients,
    )

def records_to_piecewise_polynomial(
    records: typing.Sequence[Record],
    begin: datetime.datetime,
    end: datetime.datetime,
    window: int,
    degree: int,
) -> PiecewisePolynomial:
    assert begin.tzinfo == datetime.timezone.utc
    assert end.tzinfo == datetime.timezone.utc
    assert begin < end
    if begin < records[window].time:
        raise Exception(
            f"begin ({begin}) is too close to the first record to interpolate"
        )
    if end >= records[-window].time:
        raise Exception(f"end ({end}) is too close to the last record to interpolate")
    begin_record_index = 0
    for index, record in enumerate(records):
        if begin < record.time:
            begin_record_index = index - 1
            break
    end_record_index = len(records)
    for index, record in enumerate(reversed(records)):
        if end >= record.time:
            end_record_index = len(records) - 1 - (index - 1)
            break
    return narrowed_records_to_piecewise_polynomial(
        records=list(
            itertools.islice(
                records, begin_record_index - window, end_record_index + window
            )
        ),
        window=window,
        degree=degree,
    )

def load(
    id: Id,
    begin: datetime.datetime,
    end: datetime.datetime,
    download_directory: typing.Union[str, bytes, pathlib.Path],
    window: int,
    force_download: bool = False,
) -> typing.Sequence[Record]:
    sp3_id: bytes
    if isinstance(id, Sp3Id):
        sp3_id = id.value
    elif isinstance(id, NoradId):
        sp3_id = satellite.norad_to_satellite[id.value].sp3
    else:
        raise Exception(f"unsupported id type {id.__class__}")
    assert begin.tzinfo == datetime.timezone.utc
    assert end.tzinfo == datetime.timezone.utc
    assert begin < end
    if isinstance(download_directory, bytes):
        download_directory = pathlib.Path(download_directory.decode())
    elif isinstance(download_directory, str):
        download_directory = pathlib.Path(download_directory)
    provider_found = False
    records: collections.deque[Record] = collections.deque()
    for candidate_provider in provider.find_providers_of(sp3_id):
        offset = 0.0
        begin_covered = False
        end_covered = False
        while True:
            try:
                product = Product.from_file(
                    candidate_provider.download(
                        time=candidate_provider.time_system.offset_seconds(
                            begin, offset
                        ),
                        download_directory=download_directory,
                        force=force_download,
                    )
                )
                candidate_satellite = product.satellite_with_id(sp3_id)
                if offset < 0.0:
                    for record in candidate_satellite.records[::-1]:
                        if len(records) == 0 or record.time < records[0].time:
                            records.appendleft(record)
                else:
                    for record in candidate_satellite.records:
                        if len(records) == 0 or record.time > records[-1].time:
                            records.append(record)
                if len(records) > window:
                    begin_covered = begin >= records[window].time
                    end_covered = end < records[-window].time
                    if begin_covered:
                        if end_covered:
                            provider_found = True
                            break
                        else:
                            offset = max(offset, 0.0) + candidate_provider.duration
                    else:
                        offset = min(offset, 0.0) - candidate_provider.duration
            except requests.exceptions.HTTPError as error:
                if error.response.status_code == 404:
                    logging.warning(f'"{error.request.url}" returned error 404')
                    break
                raise error
            except LookupError as error:
                if error.args[0] == sp3_id:
                    break
                raise error
        if provider_found:
            break
    if not provider_found:
        raise Exception(f'no suitable SP3 provider for "{sp3_id.decode()}"')
    return records

def obstime_to_begin_and_end(obstime: astropy.time.Time):
    assert obstime.scale == "utc"
    begin: datetime.datetime = obstime.min().to_datetime(timezone=datetime.timezone.utc)  # type: ignore
    end: datetime.datetime = obstime.max().to_datetime(timezone=datetime.timezone.utc)  # type: ignore
    while begin == end:
        end += datetime.timedelta(microseconds=1)
    return begin, end

def itrs(
    id: Id,
    obstime: astropy.time.Time,
    download_directory: typing.Union[str, bytes, pathlib.Path],
    window: int = 5,
    degree: int = 10,
) -> astropy.coordinates.ITRS:
    begin, end = obstime_to_begin_and_end(obstime)
    return records_to_piecewise_polynomial(
        records=load(
            id=id,
            begin=begin,
            end=end,
            window=window,
            download_directory=download_directory,
            force_download=False,
        ),
        begin=begin,
        end=end,
        window=window,
        degree=degree,
    )(obstime)

def altaz(
    id: Id,
    obstime: astropy.time.Time,
    location: astropy.coordinates.EarthLocation,
    pressure: astropy.units.Quantity,
    temperature: astropy.units.Quantity,
    relative_humidity: astropy.units.Quantity,
    obswl: astropy.units.Quantity,
    download_directory: typing.Union[str, bytes, pathlib.Path],
    window: int = 5,
    degree: int = 10,
) -> astropy.coordinates.AltAz:
    begin, end = obstime_to_begin_and_end(obstime)
    piecewise_polynomial = records_to_piecewise_polynomial(
        records=load(
            id=id,
            begin=begin - datetime.timedelta(seconds=2.0),
            end=end,
            window=window,
            download_directory=download_directory,
            force_download=False,
        ),
        begin=begin,
        end=end,
        window=window,
        degree=degree,
    )
    c = astropy.constants.c.to("m/s").value  # type: ignore
    location_itrs = location.get_itrs().cartesian.get_xyz()

    def light_time_correction_error(
        scalar_light_time_correction: float, scalar_obstime: astropy.time.Time
    ):
        return (
            numpy.linalg.norm(
                (
                    piecewise_polynomial(
                        scalar_obstime
                        - scalar_light_time_correction * astropy.units.Unit("s")
                    ).cartesian.get_xyz()
                    - location_itrs
                )
                .to("m")
                .value
            )
            - scalar_light_time_correction * c
        )

    light_time_correction = numpy.zeros(len(obstime))
    for index, scalar_obstime in enumerate(obstime):
        light_time_correction[index] = scipy.optimize.root_scalar(
            f=functools.partial(
                light_time_correction_error, scalar_obstime=scalar_obstime  # type: ignore
            ),
            method="brentq",
            bracket=[0.0, 2.0],
        ).root
    corrected_obstime = obstime - light_time_correction * astropy.units.Unit("s")
    itrs = piecewise_polynomial(corrected_obstime)
    itrs_vector = itrs.cartesian.get_xyz().transpose() - location_itrs
    cirs_non_topographic = astropy.coordinates.ITRS(
        x=itrs_vector[:, 0],
        y=itrs_vector[:, 1],
        z=itrs_vector[:, 2],
        v_x=itrs.v_x,
        v_y=itrs.v_y,
        v_z=itrs.v_z,
        obstime=obstime,
    ).transform_to(astropy.coordinates.CIRS(obstime=obstime))
    cirs = astropy.coordinates.CIRS(
        ra=cirs_non_topographic.ra,
        dec=cirs_non_topographic.dec,
        distance=cirs_non_topographic.distance,
        pm_ra_cosdec=cirs_non_topographic.pm_ra_cosdec,
        pm_dec=cirs_non_topographic.pm_dec,
        radial_velocity=cirs_non_topographic.radial_velocity,
        obstime=obstime,
        location=location,
    )
    altaz_vacuo = cirs.transform_to(
        astropy.coordinates.AltAz(obstime=obstime, location=location)
    )
    refraction_a, refraction_b = erfa.refco(
        phpa=pressure.to("hPa").value,
        tc=temperature.to("deg_C").value,
        rh=relative_humidity.value,
        wl=obswl.to("um").value,
    )
    distance_to_zenith = numpy.pi / 2 - altaz_vacuo.alt.to("rad").value
    tan_distance_to_zenith = numpy.tan(distance_to_zenith)
    alt_refraction_correction = (
        refraction_a * tan_distance_to_zenith
        + refraction_b * (tan_distance_to_zenith**3)
    ) * astropy.units.Unit("rad")
    # derivative of a tan(π / 2 - alt) + b tan(π / 2 - alt)³
    pm_alt_refraction_correction = (
        -altaz_vacuo.pm_alt
        * (numpy.cos(distance_to_zenith) ** -2)
        * (refraction_a + 3 * refraction_b * tan_distance_to_zenith**2)
    )
    return astropy.coordinates.AltAz(
        alt=altaz_vacuo.alt + alt_refraction_correction,
        az=altaz_vacuo.az,
        pm_az_cosalt=altaz_vacuo.pm_az_cosalt,
        pm_alt=altaz_vacuo.pm_alt + pm_alt_refraction_correction,
        obstime=obstime,
        location=location,
        pressure=pressure,
        temperature=temperature,
        relative_humidity=relative_humidity,
        obswl=0.8 * astropy.units.Unit("um"),
    )

def altaz_standard_atmosphere(
    id: Id,
    obstime: astropy.time.Time,
    location: astropy.coordinates.EarthLocation,
    download_directory: typing.Union[str, bytes, pathlib.Path],
    temperature: astropy.units.Quantity = 20.0 * astropy.units.Unit("deg_C"),  # type: ignore
    relative_humidity: astropy.units.Quantity = 0.0 * astropy.units.dimensionless_unscaled,  # type: ignore
    obswl: astropy.units.Quantity = 0.8 * astropy.units.Unit("um"),  # type: ignore
    window: int = 5,
    degree: int = 10,
):
    # https://en.wikipedia.org/wiki/Barometric_formula
    p0 = 101325 * astropy.units.Unit("Pa")
    l0 = -0.0065 * astropy.units.Unit("K/m")
    t0 = 288.15 * astropy.units.Unit("K")
    g0 = 9.80665 * astropy.units.Unit("m/s2")
    m = 0.0289644 * astropy.units.Unit("kg")
    rstar = 8.3144598 * astropy.units.Unit("J/K")
    pressure = p0 * ((1.0 + (l0 * location.height / t0)) ** ((-g0 * m) / (rstar * l0)))  # type: ignore
    return altaz(
        id=id,
        obstime=obstime,
        location=location,
        pressure=pressure,
        temperature=temperature,
        relative_humidity=relative_humidity,
        obswl=obswl,
        download_directory=download_directory,
        window=window,
        degree=degree,
    )

import numpy,import scipy,import astropy


'''
This algorithm is a IA-RL implementation on off-policy TD3 algorithm, to check the original IA-RL algorithm
you can refer to https://arxiv.org/abs/1811.06187.
Since it is a baseline algorithm, the descriptions are mostly omitted, please visit the HUGTD3.py for more implementation details
'''

import pickle
import numpy as np

import torch
import torch.nn as nn

from TD3_based_DRL.priority_replay import Memory
from TD3_based_DRL.network_model import Actor,Critic
from TD3_based_DRL.util import hard_update, soft_update

seed = 2
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

MEMORY_CAPACITY = 38400
BATCH_SIZE = 128
GAMMA = 0.95
LR_C = 0.0005
LR_A = 0.0002
LR_I = 0.01
TAU = 0.001
POLICY_NOSIE = 0.2
POLICY_FREQ = 1
NOISE_CLIP = 0.5

class DRL:

    def __init__(self, action_dim, state_dim, LR_C = LR_C, LR_A = LR_A):

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

        self.state_dim = state_dim[0] * state_dim[1]
        self.state_dim_width = state_dim[0]
        self.state_dim_height = state_dim[1]
        self.action_dim = action_dim
        self.batch_size = BATCH_SIZE
        self.gamma = GAMMA
        self.tau = TAU
        self.policy_noise = POLICY_NOSIE
        self.noise_clip = NOISE_CLIP
        self.policy_freq = POLICY_FREQ
        self.itera = 0

        self.pointer = 0
        self.memory = Memory(MEMORY_CAPACITY)

        self.actor = Actor(self.state_dim,self.action_dim).to(self.device)
        self.actor_target = Actor(self.state_dim,self.action_dim).to(self.device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),LR_A)

        self.critic = Critic(self.state_dim,self.action_dim).to(self.device)
        self.critic_target = Critic(self.state_dim,self.action_dim).to(self.device)
        self.critic_optimizers = torch.optim.Adam(self.critic.parameters(),LR_C)

        hard_update(self.actor_target,self.actor)
        hard_update(self.critic_target,self.critic)

    def learn(self, batch_size = BATCH_SIZE, epoch=0):

        ## batched state, batched action, batched action from expert, batched intervention signal, batched reward, batched next state
        bs, ba, ba_e, bi, br, bs_, tree_idx, ISweight = self.retrive(batch_size)
        bs = torch.tensor(bs, dtype=torch.float).reshape(batch_size, self.state_dim_height, self.state_dim_width).to(self.device)
        ba = torch.tensor(ba, dtype=torch.float).to(self.device).to(self.device)
        ba_e = torch.tensor(ba_e, dtype=torch.float).to(self.device).to(self.device)
        br = torch.tensor(br, dtype=torch.float).to(self.device).to(self.device)
        bs_ = torch.tensor(bs_, dtype=torch.float).reshape(batch_size, self.state_dim_height, self.state_dim_width).to(self.device)

        # initialize the loss variables
        loss_c, loss_a = 0, 0

        ## calculate the predicted values of the critic
        with torch.no_grad():
            noise1 = (torch.randn_like(ba) * self.policy_noise).clamp(0, 1)
            a_ = (self.actor_target(bs_).detach() + noise1).clamp(0, 1)
            target_q1, target_q2 = self.critic_target([bs_,a_])
            target_q1 = target_q1.detach()
            target_q2 = target_q2.detach()
            target_q = torch.min(target_q1,target_q2)
            y_expected = br + self.gamma * target_q
        y_predicted1, y_predicted2 = self.critic.forward([bs,ba])
        errors = y_expected - y_predicted1

        ## update the critic
        critic_loss = nn.MSELoss()
        loss_critic = critic_loss(y_predicted1,y_expected)+critic_loss(y_predicted2,y_expected)
        self.critic_optimizers.zero_grad()
        loss_critic.backward()
        self.critic_optimizers.step()

        ## update the actor
        if self.itera % self.policy_freq == 0:

            index1,_ = np.where(bi==0)
            index2,_ = np.where(bi==1)
            bs1,_,_,_=bs[index1],ba[index1],br[index1],bs_[index1]
            bs2,ba2,_,_=bs[index2],ba[index2],br[index2],bs_[index2]

            if bs2.size(0) != 0:
                if bs1.size(0) != 0:
                    bs1 = torch.reshape(bs1,(len(bs1), self.state_dim_height, self.state_dim_width))
                    bs2 = torch.reshape(bs2,(len(bs2), self.state_dim_height, self.state_dim_width))
                    pred_a1 = self.actor.forward(bs1)
                    pred_a2 = self.actor.forward(bs2)
                    loss_actor1 = (-self.critic.forward([bs1,pred_a1])[0])
                    ## fixed weight for human guidance actions
                    loss_actor2 = 3 * ((pred_a2 - ba2)**2)
                    loss_actor = torch.cat((loss_actor1,loss_actor2),0).mean()
                else:
                    pred_a = self.actor.forward(bs)
                    loss_actor = 3*((pred_a - ba)**2)
                    loss_actor = loss_actor.mean()
            else:
                pred_a = self.actor.forward(bs)
                loss_actor = (-self.critic.forward([bs,pred_a])[0]).mean()

            self.actor_optimizer.zero_grad()
            loss_actor.backward()
            self.actor_optimizer.step()

            soft_update(self.actor_target,self.actor,self.tau)
            soft_update(self.critic_target,self.critic,self.tau)

            loss_a = loss_actor.mean().item()

        loss_c = loss_critic.mean().item()

        self.itera += 1

        self.memory.batch_update(tree_idx, abs(errors.detach().cpu().numpy()) )

        return loss_c, loss_a

    def choose_action(self,state):

        state = torch.tensor(state,dtype=torch.float).reshape(self.state_dim_height, self.state_dim_width).to(self.device)
        state = state.unsqueeze(0)

        action = self.actor.forward(state).detach()
        action = action.squeeze(0).cpu().numpy()
        action = np.clip(action,-1, 1)

        return action

    def store_transition(self, s, a, a_e, i, r, s_):
        transition = np.hstack((s, a, a_e, i, r, s_))
        self.memory.store(transition)
        self.pointer += 1

    def retrive(self, batch_size):
        tree_index, bt, ISWeight = self.memory.sample(batch_size)
        bs = bt[:, :self.state_dim]
        ba = bt[:, self.state_dim: self.state_dim + self.action_dim]
        ba_e = bt[:, self.state_dim + self.action_dim: self.state_dim + self.action_dim + self.action_dim]
        bi = bt[:, -self.state_dim - 2: -self.state_dim - 1]
        br = bt[:, -self.state_dim - 1: -self.state_dim]
        bs_ = bt[:, -self.state_dim:]

        return bs, ba, ba_e, bi, br, bs_, tree_index, ISWeight

    def memory_save(self):

        per = open("memory_IARL.pkl", 'wb')
        str = pickle.dumps(self.memory)
        per.write(str)
        per.close()

    def memory_load(self):

        with open("memory_IARL.pkl",'rb') as file:
            self.memory  = pickle.loads(file.read())

    def load_model(self, output):
        if output is None: return
        self.actor.load_state_dict(torch.load('{}/actor.pkl'.format(output)))
        self.critic.load_state_dict(torch.load('{}/critic.pkl'.format(output)))

    def save_model(self, output):
        torch.save(self.actor.state_dict(), '{}/actor.pkl'.format(output))
        torch.save(self.critic.state_dict(), '{}/critic.pkl'.format(output))

    def save(self, log_dir, epoch):
        state = {'actor':self.actor.state_dict(), 'actor_target':self.actor_target.state_dict(),
                 'actor_optimizer':self.actor_optimizer.state_dict(),
                 'critic':self.critic.state_dict(), 'critic_target':self.critic_target.state_dict(),
                 'critic_optimizers':self.critic_optimizers.state_dict(),
                 'epoch':epoch}
        torch.save(state, log_dir)

    def load(self, log_dir):
        checkpoint = torch.load(log_dir)
        self.actor.load_state_dict(checkpoint['actor'])
        self.actor_target.load_state_dict(checkpoint['actor_target'])
        self.actor_optimizer.load_state_dict(checkpoint['actor_optimizer'])
        self.critic.load_state_dict(checkpoint['critic'])
        self.critic_target.load_state_dict(checkpoint['critic_target'])
        self.critic_optimizers.load_state_dict(checkpoint['critic_optimizers'])


#!/usr/bin/env python3

# FIXME:  Port this pipeline to 2.3.

# Copyright (c) 2015-2018 by the parties listed in the AUTHORS file.
# All rights reserved.  Use of this source code is governed by
# a BSD-style license that can be found in the LICENSE file.

from toast.mpi import MPI, finalize

import os
import sys
import time

import re
import argparse
import traceback

import pickle

import numpy as np

import toast
import toast.tod as tt
import toast.map as tm
import toast.todmap as ttm

import toast.qarray as qa
import toast.timing as timing

from toast.vis import set_backend

if tt.tidas_available:
    from toast.tod import tidas as tds

if tt.spt3g_available:
    from toast.tod import spt3g as s3g

def elapsed(mcomm, start, msg):
    mcomm.barrier()
    stop = MPI.Wtime()
    dur = stop - start
    if mcomm.rank == 0:
        print("{}: {:.3f} s".format(msg, dur), flush=True)
    return stop

def main():

    if MPI.COMM_WORLD.rank == 0:
        print("Running with {} processes".format(MPI.COMM_WORLD.size), flush=True)

    global_start = MPI.Wtime()

    parser = argparse.ArgumentParser(
        description="Read existing data and make a simple map.",
        fromfile_prefix_chars="@",
    )

    parser.add_argument(
        "--groupsize",
        required=False,
        type=int,
        default=0,
        help="size of processor groups used to distribute " "observations",
    )

    parser.add_argument(
        "--hwprpm",
        required=False,
        type=float,
        default=0.0,
        help="The rate (in RPM) of the HWP rotation",
    )

    parser.add_argument(
        "--samplerate",
        required=False,
        default=100.0,
        type=np.float,
        help="Detector sample rate (Hz)",
    )

    parser.add_argument(
        "--outdir", required=False, default="out", help="Output directory"
    )

    parser.add_argument(
        "--nside", required=False, type=int, default=64, help="Healpix NSIDE"
    )

    parser.add_argument(
        "--subnside",
        required=False,
        type=int,
        default=8,
        help="Distributed pixel sub-map NSIDE",
    )

    parser.add_argument(
        "--coord", required=False, default="E", help="Sky coordinate system [C,E,G]"
    )

    parser.add_argument(
        "--baseline",
        required=False,
        type=float,
        default=60.0,
        help="Destriping baseline length (seconds)",
    )

    parser.add_argument(
        "--noisefilter",
        required=False,
        default=False,
        action="store_true",
        help="Destripe with the noise filter enabled",
    )

    parser.add_argument(
        "--madam",
        required=False,
        default=False,
        action="store_true",
        help="If specified, use libmadam for map-making",
    )

    parser.add_argument(
        "--madampar", required=False, default=None, help="Madam parameter file"
    )

    parser.add_argument(
        "--polyorder",
        required=False,
        type=int,
        help="Polynomial order for the polyfilter",
    )

    parser.add_argument(
        "--wbin_ground",
        required=False,
        type=float,
        help="Ground template bin width [degrees]",
    )

    parser.add_argument(
        "--flush",
        required=False,
        default=False,
        action="store_true",
        help="Flush every print statement.",
    )

    parser.add_argument(
        "--tidas", required=False, default=None, help="Input TIDAS volume"
    )

    parser.add_argument(
        "--tidas_detgroup", required=False, default=None, help="TIDAS detector group"
    )

    parser.add_argument(
        "--spt3g", required=False, default=None, help="Input SPT3G data directory"
    )

    parser.add_argument(
        "--spt3g_prefix",
        required=False,
        default=None,
        help="SPT3G data frame file prefix",
    )

    parser.add_argument(
        "--common_flag_mask",
        required=False,
        default=0,
        type=np.uint8,
        help="Common flag mask",
    )

    parser.add_argument(
        "--debug",
        required=False,
        default=False,
        action="store_true",
        help="Write data distribution info and focalplane plot",
    )

    args = timing.add_arguments_and_parse(parser, timing.FILE(noquotes=True))
    # args = parser.parse_args(sys.argv)

    autotimer = timing.auto_timer("@{}".format(timing.FILE()))

    if (args.tidas is not None) and (args.spt3g is not None):
        raise RuntimeError("Cannot read two datasets!")

    if (args.tidas is None) and (args.spt3g is None):
        raise RuntimeError("No dataset specified!")

    if args.tidas is not None:
        if not tt.tidas_available:
            raise RuntimeError("TIDAS not found- cannot load")

    if args.spt3g is not None:
        if not tt.spt3g_available:
            raise RuntimeError("SPT3G not found- cannot load")

    groupsize = args.groupsize
    if groupsize == 0:
        groupsize = MPI.COMM_WORLD.size

    # Pixelization

    nside = args.nside
    npix = 12 * args.nside * args.nside
    subnside = args.subnside
    if subnside > nside:
        subnside = nside
    subnpix = 12 * subnside * subnside

    # This is the 2-level toast communicator.

    if MPI.COMM_WORLD.size % groupsize != 0:
        if MPI.COMM_WORLD.rank == 0:
            print(
                "WARNING:  process groupsize does not evenly divide into "
                "total number of processes",
                flush=True,
            )
    comm = toast.Comm(world=MPI.COMM_WORLD, groupsize=groupsize)

    # Create output directory

    mtime = MPI.Wtime()

    if comm.comm_world.rank == 0:
        if not os.path.isdir(args.outdir):
            os.makedirs(args.outdir)

    mtime = elapsed(comm.comm_world, mtime, "Creating output directory")

    # The distributed timestream data

    data = None

    if args.tidas is not None:
        if args.tidas_detgroup is None:
            raise RuntimeError("you must specify the detector group")
        data = tds.load_tidas(
            comm,
            comm.group_size,
            args.tidas,
            "r",
            args.tidas_detgroup,
            tds.TODTidas,
            group_dets=args.tidas_detgroup,
            distintervals="chunks",
        )

    if args.spt3g is not None:
        if args.spt3g_prefix is None:
            raise RuntimeError("you must specify the frame file prefix")
        data = s3g.load_spt3g(
            comm,
            comm.group_size,
            args.spt3g,
            args.spt3g_prefix,
            s3g.obsweight_spt3g,
            s3g.TOD3G,
        )

    mtime = elapsed(comm.comm_world, mtime, "Distribute data")

    # In debug mode, print out data distribution information

    if args.debug:
        handle = None
        if comm.comm_world.rank == 0:
            handle = open("{}_distdata.txt".format(args.outdir), "w")
        data.info(handle)
        if comm.comm_world.rank == 0:
            handle.close()
        mtime = elapsed(comm.comm_world, mtime, "Dumping debug data distribution")
        if comm.comm_world.rank == 0:
            outfile = "{}_focalplane.png".format(args.outdir)
            set_backend()
            # Just plot the dets from the first TOD
            temptod = data.obs[0]["tod"]
            # FIXME: change this once we store det info in the metadata.
            dfwhm = {x: 10.0 for x in temptod.detectors}
            tt.plot_focalplane(temptod.detoffset(), 10.0, 10.0, outfile, fwhm=dfwhm)
        comm.comm_world.barrier()
        mtime = elapsed(comm.comm_world, mtime, "Plotting debug focalplane")

    # Compute pointing matrix

    pointing = tt.OpPointingHpix(
        nside=args.nside, nest=True, mode="IQU", hwprpm=args.hwprpm
    )
    pointing.exec(data)

    mtime = elapsed(comm.comm_world, mtime, "Expand pointing")

    # Mapmaking.

    # FIXME:  We potentially have a different noise model for every
    # observation.  We need to have both spt3g and tidas format Noise
    # classes which read the information from disk.  Then the mapmaking
    # operators need to get these noise weights from each observation.
    detweights = {d: 1.0 for d in data.obs[0]["tod"].detectors}

    if not args.madam:
        if comm.comm_world.rank == 0:
            print("Not using Madam, will only make a binned map!", flush=True)

        # Filter data if desired

        if args.polyorder:
            polyfilter = tt.OpPolyFilter(
                order=args.polyorder, common_flag_mask=args.common_flag_mask
            )
            polyfilter.exec(data)
            mtime = elapsed(comm.comm_world, mtime, "Polynomial filtering")

        if args.wbin_ground:
            groundfilter = tt.OpGroundFilter(
                wbin=args.wbin_ground, common_flag_mask=args.common_flag_mask
            )
            groundfilter.exec(data)
            mtime = elapsed(comm.comm_world, mtime, "Ground template filtering")

        # Compute pixel space distribution

        lc = tm.OpLocalPixels()
        localpix = lc.exec(data)
        if localpix is None:
            raise RuntimeError(
                "Process {} has no hit pixels. Perhaps there are fewer "
                "detectors than processes in the group?".format(comm.comm_world.rank)
            )
        localsm = np.unique(np.floor_divide(localpix, subnpix))
        mtime = elapsed(comm.comm_world, mtime, "Compute local submaps")

        # construct distributed maps to store the covariance,
        # noise weighted map, and hits

        mtime = MPI.Wtime()
        invnpp = tm.DistPixels(
            comm=comm.comm_world,
            size=npix,
            nnz=6,
            dtype=np.float64,
            submap=subnpix,
            local=localsm,
        )
        hits = tm.DistPixels(
            comm=comm.comm_world,
            size=npix,
            nnz=1,
            dtype=np.int64,
            submap=subnpix,
            local=localsm,
        )
        zmap = tm.DistPixels(
            comm=comm.comm_world,
            size=npix,
            nnz=3,
            dtype=np.float64,
            submap=subnpix,
            local=localsm,
        )

        # compute the hits and covariance.

        invnpp.data.fill(0.0)
        hits.data.fill(0)

        build_invnpp = tm.OpAccumDiag(
            detweights=detweights,
            invnpp=invnpp,
            hits=hits,
            common_flag_mask=args.common_flag_mask,
        )
        build_invnpp.exec(data)

        invnpp.allreduce()
        hits.allreduce()
        mtime = elapsed(comm.comm_world, mtime, "Building hits and N_pp^-1")

        hits.write_healpix_fits("{}_hits.fits".format(args.outdir))
        invnpp.write_healpix_fits("{}_invnpp.fits".format(args.outdir))
        mtime = elapsed(comm.comm_world, mtime, "Writing hits and N_pp^-1")

        # invert it
        tm.covariance_invert(invnpp, 1.0e-3)
        mtime = elapsed(comm.comm_world, mtime, "Inverting N_pp^-1")

        invnpp.write_healpix_fits("{}_npp.fits".format(args.outdir))
        mtime = elapsed(comm.comm_world, mtime, "Writing N_pp")

        zmap.data.fill(0.0)
        build_zmap = tm.OpAccumDiag(
            zmap=zmap, detweights=detweights, common_flag_mask=args.common_flag_mask
        )
        build_zmap.exec(data)
        zmap.allreduce()
        mtime = elapsed(comm.comm_world, mtime, "Building noise weighted map")

        tm.covariance_apply(invnpp, zmap)
        mtime = elapsed(comm.comm_world, mtime, "Computing binned map")

        zmap.write_healpix_fits(os.path.join(args.outdir, "binned.fits"))
        mtime = elapsed(comm.comm_world, mtime, "Writing binned map")

    else:
        # Set up MADAM map making.

        pars = {}
        pars["temperature_only"] = "F"
        pars["force_pol"] = "T"
        pars["kfirst"] = "T"
        pars["concatenate_messages"] = "T"
        pars["write_map"] = "T"
        pars["write_binmap"] = "T"
        pars["write_matrix"] = "T"
        pars["write_wcov"] = "T"
        pars["write_hits"] = "T"
        pars["nside_cross"] = nside // 2
        pars["nside_submap"] = subnside

        if args.madampar is not None:
            pat = re.compile(r"\s*(\S+)\s*=\s*(\S+(\s+\S+)*)\s*")
            comment = re.compile(r"^#.*")
            with open(args.madampar, "r") as f:
                for line in f:
                    if comment.match(line) is None:
                        result = pat.match(line)
                        if result is not None:
                            key, value = result.group(1), result.group(2)
                            pars[key] = value

        pars["base_first"] = args.baseline
        pars["nside_map"] = nside
        if args.noisefilter:
            pars["kfilter"] = "T"
        else:
            pars["kfilter"] = "F"
        pars["fsample"] = args.samplerate

        madam = tm.OpMadam(
            params=pars, detweights=detweights, common_flag_mask=args.common_flag_mask
        )
        madam.exec(data)
        mtime = elapsed(comm.comm_world, mtime, "Madam mapmaking")

    comm.comm_world.barrier()
    stop = MPI.Wtime()
    dur = stop - global_start
    if comm.comm_world.rank == 0:
        print("Total Time:  {:.2f} seconds".format(dur), flush=True)
    return

if __name__ == "__main__":
    try:
        main()
        tman = timing.timing_manager()
        tman.report()
    except:
        exc_type, exc_value, exc_traceback = sys.exc_info()
        lines = traceback.format_exception(exc_type, exc_value, exc_traceback)
        lines = ["Proc {}: {}".format(MPI.COMM_WORLD.rank, x) for x in lines]
        print("".join(lines), flush=True)
        toast.raise_error(6)  # typical error code for SIGABRT
        MPI.COMM_WORLD.Abort(6)
    finalize()


#!/usr/bin/env python

import numpy as np
import Pdata
from collections import OrderedDict
import gnc
import pandas as pa

"""
This module stores functions only meaningful to ORCHIDEE
"""

veget4type=['bareland','forest','grass','crop']
veget5type=['bareland','forest','grass','pasture','crop']
forest4type = ['BrEv','BrDe','NeEv','NeDe']
forest4type_longname = ['BrEv','BrDe','NeEv','NeDe']

def VEGET_MAX_index_list(typenum=4):
    """
    Return the np.s_ index to slice the vegetmax related variables
        into different classes.

    Parameters:
    -----------
    typenum:
        4: recast into bareland,forest,grass,and crop.
        5: recast into bareland,forest,grass,pasture,and crop.
        'forest': extract only the fractions for forest into four types:
            BrEv,BrDe,NeEv,NeDe
        'Jinfeng17':'bareland','forest','grass','crop','pasture','rangeland'
    """
    if typenum == 4:
        index_list = [np.s_[...,0:1,:,:],
                      np.s_[...,1:9,:,:],
                      np.s_[...,9:11,:,:],
                      np.s_[...,11:13,:,:]
                     ]
        namelist = ['bareland','forest','grass','crop']

    elif typenum == 'forest':
        index_list = [np.s_[...,[1,4],:,:],
                      np.s_[...,[2,5,7],:,:],
                      np.s_[...,[3,6],:,:],
                      np.s_[...,[8],:,:]
                     ]
        namelist = ['BrEv','BrDe','NeEv','NeDe']

    elif typenum == 5:
        index_list = [np.s_[...,0:1,:,:],
                      np.s_[...,1:9,:,:],
                      np.s_[...,[9,11],:,:],
                      np.s_[...,[10,12],:,:],
                      np.s_[...,[13,14],:,:]
                     ]
        namelist = ['bareland','forest','grass','pasture','crop']

    elif typenum == 'Jinfeng17':
        index_list = [np.s_[...,0:1,:,:],
                      np.s_[...,1:9,:,:],
                      np.s_[...,[9,10],:,:],
                      np.s_[...,[11,12],:,:],
                      np.s_[...,[13,14],:,:],
                      np.s_[...,[15,16],:,:]
                     ]
        namelist = ['bareland','forest','grass','crop','pasture','rangeland']

    else:
        raise ValueError("typenum not correct ! ")
    return (index_list,namelist)

def VEGET_MAX_recast(arr,typenum=4,return_dic=False):
    """
    change the VEGET_MAX variable into another several types.

    Parameters:
    -----------
    arr: input VEGET_MAX or veget ndarray.
    typenum:
        4: recast into bareland,forest,grass,and agriculture.
        5: recast into bareland,forest,grass,pasture and crop.
        forest: extract only the fractions for forest into four types:
            BrEv,BrDe,NeEv,NeDe
    return_dic: boolean, True to return a dictionary with the keys as the types.
    """
    if arr.ndim < 3:
        raise ValueError("input array dimension less than 3!")
    else:
        if arr.shape[-3] not in [13,15,17]:
            raise ValueError("the third last dim is not PFT!")
        else:
            index_list,namelist = VEGET_MAX_index_list(typenum)
            arrlist = [np.ma.sum(arr[ind],axis=-3)[...,np.newaxis,:,:]
                        for ind in index_list]
            if not return_dic:
                return np.ma.concatenate(arrlist,axis=-3)
            else:
                arrlist = [arr[...,0,:,:] for arr in arrlist]
                return OrderedDict(zip(namelist,arrlist))

def dataframe_from_stomate(filepattern,largefile=True,multifile=True,
                           dgvmadj=False,spamask=None,
                           veget_npindex=np.s_[:],areaind=np.s_[:],
                           out_timestep='annual',version=1,
                           replace_nan=False):
    """
    Parameters:
    -----------
    filepattern: could be a single filename, or a file pattern
    out_timestep: the timestep of output file, used to provide information
        to properly scale the variable values, could be 'annual' or 'daily'.
        when 'annual', flux_scale_factor = 365 will be used.

    dgvmadj: use DGVM adjustment, in this case tBIOMASS rathern than TOTAL_M
        is used.
    veget_npindex: passed to the function of get_pftsum:
        1. could be used to restrict for example the PFT
        weighted average only among natural PFTs by setting
        veget_npindex=np.s_[:,0:11,:,:]. It will be used to slice
        VEGET_MAX variable.
        2. could also be used to slice only for some subgrid
        of the whole grid, eg., veget_npindex=np.s_[...,140:300,140:290].

    Notes:
    ------
    1. This function could handle automatically the case of a single-point
       file or a regional file. When a single-point file (pattern) is given,
       PFT-weighted carbon density will be used rather than the total C over
       the spatial area.
    """
    gnc_sto = gnc.Ncdata(filepattern,largefile=largefile,multifile=multifile,
                         replace_nan=replace_nan)

    if version == 1:
        # list all pools and fluxes
        list_flux_pft = ['GPP','NPP','HET_RESP','CO2_FIRE','CO2FLUX','CO2_TAKEN']
        list_flux_pftsum = ['CONVFLUX','CFLUX_PROD10','CFLUX_PROD100','HARVEST_ABOVE']
        list_flux = list_flux_pft+list_flux_pftsum

        list_pool = ['TOTAL_M','TOTAL_SOIL_CARB']
        list_all = list_flux_pft+list_flux_pftsum+list_pool
        nlist_var = [list_flux_pft, list_flux_pftsum, list_pool]

        for varlist in nlist_var:
            gnc_sto.retrieve_variables(varlist)
            gnc_sto.get_pftsum(print_info=False,veget_npindex=veget_npindex)
            gnc_sto.remove_variables(varlist)

        #handle adjustment of different variables
        if dgvmadj:
            gnc_sto.retrieve_variables(['tGPP','tRESP_GROWTH','tRESP_MAINT','tRESP_HETERO','tCO2_FIRE'])
            gnc_sto.pftsum.__dict__['NPP'] = gnc_sto.d1.tGPP - gnc_sto.d1.tRESP_MAINT - gnc_sto.d1.tRESP_GROWTH
            gnc_sto.pftsum.__dict__['HET_RESP'] = gnc_sto.d1.tRESP_HETERO
            gnc_sto.pftsum.__dict__['CO2_FIRE'] = gnc_sto.d1.tCO2_FIRE
            gnc_sto.remove_variables(['tGPP','tRESP_GROWTH','tRESP_MAINT','tRESP_HETERO','tCO2_FIRE'])

            gnc_sto.retrieve_variables(['tBIOMASS','tLITTER','tSOILC'])
            gnc_sto.pftsum.__dict__['TOTAL_M'] = gnc_sto.d1.tBIOMASS
            gnc_sto.pftsum.__dict__['TOTAL_SOIL_CARB'] = gnc_sto.d1.tLITTER + gnc_sto.d1.tSOILC
            gnc_sto.remove_variables(['tBIOMASS','tLITTER','tSOILC'])

        # we have to treat product pool independently
        try:
            gnc_sto.retrieve_variables(['PROD10','PROD100'])
            gnc_sto.pftsum.PROD10 = gnc_sto.d1.PROD10.sum(axis=1)
            gnc_sto.pftsum.PROD100 = gnc_sto.d1.PROD100.sum(axis=1)
            gnc_sto.remove_variables(['PROD10','PROD100'])
        except KeyError:
            gnc_sto.pftsum.PROD10 = gnc_sto.pftsum.NPP * 0.
            gnc_sto.pftsum.PROD100 = gnc_sto.pftsum.NPP * 0.

        # get the spatial operation and pass them into dataframe
        if not gnc_sto._SinglePoint:
            gnc_sto.get_spa()
            dft = pa.DataFrame(gnc_sto.spasum.__dict__)
        else:
            dft = pa.DataFrame(gnc_sto.pftsum.__dict__)

        # treat the output time step
        if out_timestep == 'annual':
            flux_scale_factor = 365.
            dft['CO2FLUX'] = dft['CO2FLUX']/30.  #CO2FLUX is monthly output
        elif out_timestep == 'daily':
            flux_scale_factor = 1
        dft[list_flux] = dft[list_flux]*flux_scale_factor

        # get total carbon pool
        dft['PROD'] = dft['PROD10'] + dft['PROD100']
        dft['CarbonPool'] = dft['TOTAL_M'] + dft['TOTAL_SOIL_CARB'] + dft['PROD']

        # calcate NBP
        dft['NBP_npp'] = dft['NPP']+dft['CO2_TAKEN']-dft['CONVFLUX']-dft['CFLUX_PROD10']-dft['CFLUX_PROD100']-dft['CO2_FIRE']-dft['HARVEST_ABOVE']-dft['HET_RESP']
        dft['NBP_co2flux'] = -1*(dft['CO2FLUX']+dft['HARVEST_ABOVE']+dft['CONVFLUX']+dft['CFLUX_PROD10']+dft['CFLUX_PROD100'])

    elif version == 2:
        # list all pools and fluxes
        list_flux_pft = ['GPP','NPP','HET_RESP','CO2_FIRE','CO2FLUX','CO2_TAKEN','METHANE','RANIMAL']
        list_flux_pftsum = ['CONVFLUX_LCC','CONVFLUX_HAR','CFLUX_PROD10_LCC','CFLUX_PROD10_HAR','CFLUX_PROD100_LCC','CFLUX_PROD100_HAR','HARVEST_ABOVE']
        list_flux = list_flux_pft+list_flux_pftsum

        list_pool = ['TOTAL_M','TOTAL_SOIL_CARB','LEAF_M','SAP_M_AB','SAP_M_BE',
                     'HEART_M_AB','HEART_M_BE','ROOT_M','FRUIT_M','RESERVE_M',
                     'LITTER_STR_AB','LITTER_STR_BE','LITTER_MET_AB','LITTER_MET_BE']
        list_all = list_flux_pft+list_flux_pftsum+list_pool
        nlist_var = [list_flux_pft, list_flux_pftsum, list_pool]

        for varlist in nlist_var:
            gnc_sto.retrieve_variables(varlist,mask=spamask)
            gnc_sto.get_pftsum(print_info=False,veget_npindex=veget_npindex)
            gnc_sto.remove_variables(varlist)

        #handle adjustment of different variables
        if dgvmadj:
            if veget_npindex != np.s_[:]:
                raise ValueError("dgvmadj is not handled when veget_npindex does not include all")
            else:
                gnc_sto.retrieve_variables(['tGPP','tRESP_GROWTH','tRESP_MAINT','tRESP_HETERO','tCO2_FIRE'],mask=spamask)
                gnc_sto.pftsum.__dict__['NPP'] = gnc_sto.d1.tGPP - gnc_sto.d1.tRESP_MAINT - gnc_sto.d1.tRESP_GROWTH
                gnc_sto.pftsum.__dict__['HET_RESP'] = gnc_sto.d1.tRESP_HETERO
                gnc_sto.pftsum.__dict__['CO2_FIRE'] = gnc_sto.d1.tCO2_FIRE
                gnc_sto.remove_variables(['tGPP','tRESP_GROWTH','tRESP_MAINT','tRESP_HETERO','tCO2_FIRE'])

                gnc_sto.retrieve_variables(['tBIOMASS','tLITTER','tSOILC'],mask=spamask)
                gnc_sto.pftsum.__dict__['TOTAL_M'] = gnc_sto.d1.tBIOMASS
                gnc_sto.pftsum.__dict__['TOTAL_SOIL_CARB'] = gnc_sto.d1.tLITTER + gnc_sto.d1.tSOILC
                gnc_sto.remove_variables(['tBIOMASS','tLITTER','tSOILC'])

        # we have to treat product pool independently
        list_prod = ['PROD10_LCC','PROD10_HAR','PROD100_LCC','PROD100_HAR']
        gnc_sto.retrieve_variables(list_prod,mask=spamask)
        for var in list_prod:
            gnc_sto.pftsum.__dict__[var] = gnc_sto.d1.__dict__[var][veget_npindex].sum(axis=1)
        print gnc_sto.d1.__dict__['PROD10_LCC'][veget_npindex].shape
        print gnc_sto.d1.__dict__['PROD10_LCC'].shape
        print gnc_sto.pftsum.__dict__['PROD10_LCC'].shape
        gnc_sto.remove_variables(list_prod)

        # get the spatial operation and pass them into dataframe
        if not gnc_sto._SinglePoint:
            gnc_sto.get_spa(areaind=areaind)
            dft = pa.DataFrame(gnc_sto.spasum.__dict__)
        else:
            dft = pa.DataFrame(gnc_sto.pftsum.__dict__)

        #  2016-03-30: the shape of gnc_sto.d1.ContAreas could be
        #  (nlat,nlon) when there is no "CONTFRAC" or "NONBIOFRAC" in
        #  the history file, but could be (ntime,nlat,nlon) when they're
        #  present.

        #  # [++temporary++] treat CO2_TAKEN
        #  # In case of shifting cultivation is simulated, the CO2_TAKEN
        #  # could be big at the last day. However the veget_max is kept
        #  # the same as the old one over the year, so we have to use
        #  # last-year CO2_TAKEN multiply with the next-year veget_max.
        #  gnc_sto.retrieve_variables(['CO2_TAKEN'])
        #  co2taken_pftsum = np.ma.sum(gnc_sto.d1.CO2_TAKEN[:-1] * gnc_sto.d1.VEGET_MAX[1:],axis=1)
        #  if not gnc_sto._SinglePoint:
        #      dt = np.sum(co2taken_pftsum*gnc_sto.d1.ContAreas,axis=(1,2))
        #  else:
        #      dt = co2taken_pftsum
        #  dft['CO2_TAKEN'].iloc[:-1] = dt

        # treat the output time step
        if out_timestep == 'annual':
            flux_scale_factor = 365.
            dft['CO2FLUX'] = dft['CO2FLUX']/30.  #CO2FLUX is monthly output
        elif out_timestep == 'daily':
            flux_scale_factor = 1
        dft[list_flux] = dft[list_flux]*flux_scale_factor

        # get total carbon pool
        dft['PROD'] = dft['PROD10_LCC'] + dft['PROD10_HAR'] + dft['PROD100_LCC'] + dft['PROD100_HAR']
        dft['CarbonPool'] = dft['TOTAL_M'] + dft['TOTAL_SOIL_CARB'] + dft['PROD']
        dft['LITTER_AB'] = dft['LITTER_STR_AB'] + dft['LITTER_MET_AB']
        dft['LITTER_BE'] = dft['LITTER_MET_BE'] + dft['LITTER_STR_BE']
        dft['LITTER'] = dft['LITTER_BE'] + dft['LITTER_AB']
        dft['BIOMASS_AB'] = dft.SAP_M_AB + dft.HEART_M_AB + dft.LEAF_M + dft.FRUIT_M + dft.RESERVE_M
        dft['BIOMASS_BE'] = dft.SAP_M_BE + dft.HEART_M_BE + dft.ROOT_M

        # treat GM
        dft['RANIMAL'] = dft['RANIMAL']*1000
        dft['METHANE'] = dft['METHANE']*1000
        dft['GMsource'] = dft['RANIMAL'] + dft['METHANE']

        # treat LUC
        dft['CONVFLUX'] = dft['CONVFLUX_LCC'] + dft['CONVFLUX_HAR']
        dft['CFLUX_PROD10'] = dft['CFLUX_PROD10_LCC'] + dft['CFLUX_PROD10_HAR']
        dft['CFLUX_PROD100'] = dft['CFLUX_PROD100_LCC'] + dft['CFLUX_PROD100_HAR']
        dft['LUCsource'] = dft['CONVFLUX'] + dft['CFLUX_PROD10'] + dft['CFLUX_PROD100']

        # calcate NBP
        dft['NBP_npp'] = dft['NPP']+dft['CO2_TAKEN']-dft['CONVFLUX']-dft['CFLUX_PROD10']-dft['CFLUX_PROD100']-dft['CO2_FIRE'] \
                         -dft['HARVEST_ABOVE']-dft['HET_RESP']-dft['RANIMAL']-dft['METHANE']
        dft['NBP_co2flux'] = -1*(dft['CO2FLUX']+dft['HARVEST_ABOVE']+dft['CONVFLUX']+dft['CFLUX_PROD10']+dft['CFLUX_PROD100'])

        # litter
        dft['LITTER'] = dft[['LITTER_STR_AB','LITTER_STR_BE','LITTER_MET_AB','LITTER_MET_BE']].sum(axis=1)
        dft['LITTER_AB'] = dft[['LITTER_STR_AB','LITTER_MET_AB']].sum(axis=1)
        dft['LITTER_BE'] = dft[['LITTER_STR_BE','LITTER_MET_BE']].sum(axis=1)
        dft['SOILC'] = dft['TOTAL_SOIL_CARB'] - dft['LITTER']

    else:
        raise ValueError("Unknown version!")

    gnc_sto.close()

    return dft

def panel_summary_stomate(dic):
    """
    Return a panel of summary dataframe for stomate_history.nc using
    dataframe_from_stomate

    Parameters:
    -----------
    dic: a dictionary of (tag,filename) pairs.
    """
    pdic = OrderedDict()
    for k,f in dic.items():
        pdic[k] = dataframe_from_stomate(f)
    return pa.Panel(pdic)

def write_PFTmap(filename,data,resolution='05deg'):
    if resolution=='05deg':
        ### Build the dimensions
        ncfile = gnc.NcWrite(filename)
        ncfile.add_dim_lat(bounds = "bounds_lat",latvar=np.arange(89.75,-90,-0.5),units = "degrees_north",valid_min = -90.,valid_max = 90.,long_name = "Latitude",axis = "Y")
        ncfile.add_dim_lon(bounds = "bounds_lon",lonvar=np.arange(-179.75,180,0.5),units = "degrees_east",valid_min = -180.,valid_max = 180.,long_name = "Longitude",axis = "X")
        timedim_info = ['time_counter','time_counter','time_counter','f4',np.array([1850]),'no unit',True]
        ncfile.add_dim(timedim_info,units = "years since 0-1-1",Calendar = "gregorian",axis = "T")
        vegetdim_info = ['veget','veget','Vegetation Classes','i4',np.arange(1,16),"-",False]
        ncfile.add_dim(vegetdim_info,validmax = 15.,validmin = 1.)
        import cons
        from collections import OrderedDict
        pftdic = OrderedDict()
        for i in range(1,16):
            pftdic['PFT{0:0>2}'.format(i)] = cons.pftdic15[i]
            ### Write the vegetmax
        varinfo_value = ['maxvegetfrac',('time_counter', 'veget', 'lat', 'lon',),'f4',data]
        ncfile.add_var(varinfo_value,name = "maxvegetfrac",long_name = "Vegetation types",units = "-")
        ncfile.add_global_attributes(pftdic)
        ncfile.add_history_attr(histtxt="PFT map generated using data from FireMIP")
        ncfile.close()


# Copyright 2019-2022 Cambridge Quantum Computing
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
from typing import cast, Iterable, List, Optional, Sequence, Union, Any
from uuid import uuid4
from logging import warning

import numpy as np
from pyquil.api import (
    QuantumComputer,
    WavefunctionSimulator,
    list_quantum_computers,
    get_qc,
)
from pyquil.gates import I
from pyquil.paulis import ID, PauliSum, PauliTerm
from pyquil.quilatom import Qubit as Qubit_

from pytket.circuit import Circuit, OpType, Qubit  # type: ignore
from pytket.backends import (
    Backend,
    CircuitNotRunError,
    CircuitStatus,
    ResultHandle,
    StatusEnum,
)
from pytket.backends.backend import KwargTypes
from pytket.backends.backendinfo import BackendInfo
from pytket.backends.backendresult import BackendResult
from pytket.backends.resulthandle import _ResultIdTuple
from pytket.extensions.pyquil._metadata import __extension_version__
from pytket.passes import (  # type: ignore
    BasePass,
    EulerAngleReduction,
    CXMappingPass,
    auto_rebase_pass,
    SequencePass,
    SynthesiseTket,
    DecomposeBoxes,
    FullPeepholeOptimise,
    CliffordSimp,
    FlattenRegisters,
    SimplifyInitial,
    NaivePlacementPass,
)
from pytket.pauli import QubitPauliString  # type: ignore
from pytket.predicates import (  # type: ignore
    NoSymbolsPredicate,
    ConnectivityPredicate,
    GateSetPredicate,
    NoClassicalControlPredicate,
    NoFastFeedforwardPredicate,
    NoMidMeasurePredicate,
    DefaultRegisterPredicate,
    Predicate,
)
from pytket.extensions.pyquil.pyquil_convert import (
    process_characterisation,
    get_avg_characterisation,
    tk_to_pyquil,
)
from pytket.placement import NoiseAwarePlacement  # type: ignore
from pytket.architecture import Architecture  # type: ignore
from pytket.utils import prepare_circuit
from pytket.utils.operators import QubitPauliOperator
from pytket.utils.outcomearray import OutcomeArray

class PyQuilJobStatusUnavailable(Exception):
    """Raised when trying to retrieve unknown job status."""

    def __init__(self) -> None:
        super().__init__("The job status cannot be retrieved.")

_STATUS_MAP = {
    "done": StatusEnum.COMPLETED,
    "running": StatusEnum.RUNNING,
    "loaded": StatusEnum.SUBMITTED,
    "connected": StatusEnum.SUBMITTED,
}

def _default_q_index(q: Qubit) -> int:
    if q.reg_name != "q" or len(q.index) != 1:
        raise ValueError("Non-default qubit register")
    return int(q.index[0])

class ForestBackend(Backend):
    _supports_shots = True
    _supports_counts = True
    _supports_contextual_optimisation = True
    _persistent_handles = True
    _GATE_SET = {OpType.CZ, OpType.Rx, OpType.Rz, OpType.Measure, OpType.Barrier}

    def __init__(self, qc: QuantumComputer):
        """Backend for running circuits with the Rigetti QVM.

        :param qc: The particular QuantumComputer to use. See the pyQuil docs for more
        details.
        :type qc: QuantumComputer
        """
        super().__init__()
        self._qc: QuantumComputer = qc
        self._backend_info = self._get_backend_info(self._qc)

    @property
    def required_predicates(self) -> List[Predicate]:
        return [
            NoClassicalControlPredicate(),
            NoFastFeedforwardPredicate(),
            NoMidMeasurePredicate(),
            GateSetPredicate(self.backend_info.gate_set),
            ConnectivityPredicate(self.backend_info.architecture),
        ]

    def rebase_pass(self) -> BasePass:
        return auto_rebase_pass({OpType.CZ, OpType.Rz, OpType.Rx})

    def default_compilation_pass(self, optimisation_level: int = 1) -> BasePass:
        assert optimisation_level in range(3)
        passlist = [
            DecomposeBoxes(),
            FlattenRegisters(),
        ]
        if optimisation_level == 1:
            passlist.append(SynthesiseTket())
        elif optimisation_level == 2:
            passlist.append(FullPeepholeOptimise())
        passlist.append(
            CXMappingPass(
                self.backend_info.architecture,
                NoiseAwarePlacement(
                    self._backend_info.architecture,
                    self._backend_info.averaged_node_gate_errors,
                    self._backend_info.averaged_edge_gate_errors,
                ),
                directed_cx=False,
                delay_measures=True,
            )
        )
        passlist.append(NaivePlacementPass(self.backend_info.architecture))
        if optimisation_level == 2:
            passlist.append(CliffordSimp(False))
        if optimisation_level > 0:
            passlist.append(SynthesiseTket())
        passlist.append(self.rebase_pass())
        if optimisation_level > 0:
            passlist.extend(
                [
                    EulerAngleReduction(OpType.Rx, OpType.Rz),
                    SimplifyInitial(
                        allow_classical=False, create_all_qubits=True, xcirc=_xcirc
                    ),
                ]
            )
        return SequencePass(passlist)

    @property
    def _result_id_type(self) -> _ResultIdTuple:
        return (int, str)

    def process_circuits(
        self,
        circuits: Sequence[Circuit],
        n_shots: Union[None, int, Sequence[Optional[int]]] = None,
        valid_check: bool = True,
        **kwargs: KwargTypes,
    ) -> List[ResultHandle]:
        """
        See :py:meth:`pytket.backends.Backend.process_circuits`.
        Supported kwargs: `seed`.
        """
        circuits = list(circuits)
        n_shots_list = Backend._get_n_shots_as_list(
            n_shots, len(circuits), optional=False
        )

        if valid_check:
            self._check_all_circuits(circuits)

        postprocess = kwargs.get("postprocess", False)

        handle_list = []
        for circuit, n_shots in zip(circuits, n_shots_list):
            if postprocess:
                c0, ppcirc = prepare_circuit(circuit, allow_classical=False)
                ppcirc_rep = ppcirc.to_dict()
            else:
                c0, ppcirc_rep = circuit, None
            p, bits = tk_to_pyquil(c0, return_used_bits=True)
            p.wrap_in_numshots_loop(n_shots)
            ex = self._qc.compiler.native_quil_to_executable(p)
            qam = self._qc.qam
            qam.random_seed = kwargs.get("seed")  # type: ignore
            pyquil_handle = qam.execute(ex)
            handle = ResultHandle(uuid4().int, json.dumps(ppcirc_rep))
            measures = circuit.n_gates_of_type(OpType.Measure)
            if measures == 0:
                self._cache[handle] = {
                    "handle": pyquil_handle,
                    "c_bits": sorted(bits),
                    "result": self.empty_result(circuit, n_shots=n_shots),
                }
            else:
                self._cache[handle] = {"handle": pyquil_handle, "c_bits": sorted(bits)}
            handle_list.append(handle)
        return handle_list

    def circuit_status(self, handle: ResultHandle) -> CircuitStatus:
        """
        Return a CircuitStatus reporting the status of the circuit execution
        corresponding to the ResultHandle.

        This will throw an PyQuilJobStatusUnavailable exception if the results
        have not been retrieved yet, as pyQuil does not currently support asynchronous
        job status queries.

        :param handle: The handle to the submitted job.
        :type handle: ResultHandle
        :returns: The status of the submitted job.
        :raises PyQuilJobStatusUnavailable: Cannot retrieve job status.
        :raises CircuitNotRunError: The handle does not correspond to a valid job.
        """
        if handle in self._cache and "result" in self._cache[handle]:
            return CircuitStatus(StatusEnum.COMPLETED)
        if handle in self._cache:
            # retrieving status is not supported yet
            # see https://github.com/rigetti/pyquil/issues/1370
            raise PyQuilJobStatusUnavailable()
        raise CircuitNotRunError(handle)

    def get_result(self, handle: ResultHandle, **kwargs: KwargTypes) -> BackendResult:
        """
        See :py:meth:`pytket.backends.Backend.get_result`.
        Supported kwargs: none.
        """
        try:
            return super().get_result(handle)
        except CircuitNotRunError:
            if handle not in self._cache:
                raise CircuitNotRunError(handle)

            pyquil_handle = self._cache[handle]["handle"]
            raw_shots = self._qc.qam.get_result(pyquil_handle).readout_data["ro"]
            if raw_shots is None:
                raise ValueError("Could not read job results in memory")
            shots = OutcomeArray.from_readouts(raw_shots.tolist())
            ppcirc_rep = json.loads(cast(str, handle[1]))
            ppcirc = Circuit.from_dict(ppcirc_rep) if ppcirc_rep is not None else None
            res = BackendResult(
                shots=shots, c_bits=self._cache[handle]["c_bits"], ppcirc=ppcirc
            )
            self._cache[handle].update({"result": res})
            return res

    @property
    def backend_info(self) -> BackendInfo:
        return self._backend_info

    @classmethod
    def _get_backend_info(cls, qc: QuantumComputer) -> BackendInfo:
        char_dict: dict = process_characterisation(qc)
        arch = char_dict.get("Architecture", Architecture([]))
        node_errors = char_dict.get("NodeErrors")
        link_errors = char_dict.get("EdgeErrors")
        averaged_errors = get_avg_characterisation(char_dict)
        return BackendInfo(
            cls.__name__,
            qc.name,
            __extension_version__,
            arch,
            cls._GATE_SET,
            all_node_gate_errors=node_errors,
            all_edge_gate_errors=link_errors,
            averaged_node_gate_errors=averaged_errors["node_errors"],
            averaged_edge_gate_errors=averaged_errors["link_errors"],
        )

    @classmethod
    def available_devices(cls, **kwargs: Any) -> List[BackendInfo]:
        """
        See :py:meth:`pytket.backends.Backend.available_devices`.
        Supported kwargs: `qpus` (default true), `qvms` (default false).
        """
        if "qvms" not in kwargs:
            kwargs["qvms"] = False
        qc_name_list = list_quantum_computers(**kwargs)
        return [cls._get_backend_info(get_qc(name)) for name in qc_name_list]

class ForestStateBackend(Backend):
    _supports_state = True
    _supports_expectation = True
    _expectation_allows_nonhermitian = False
    _persistent_handles = False
    _GATE_SET = {
        OpType.X,
        OpType.Y,
        OpType.Z,
        OpType.H,
        OpType.S,
        OpType.T,
        OpType.Rx,
        OpType.Ry,
        OpType.Rz,
        OpType.CZ,
        OpType.CX,
        OpType.CCX,
        OpType.CU1,
        OpType.U1,
        OpType.SWAP,
    }

    def __init__(self) -> None:
        """Backend for running simulations on the Rigetti QVM Wavefunction Simulator."""
        super().__init__()
        self._sim = WavefunctionSimulator()

    @property
    def required_predicates(self) -> List[Predicate]:
        return [
            NoClassicalControlPredicate(),
            NoFastFeedforwardPredicate(),
            NoMidMeasurePredicate(),
            NoSymbolsPredicate(),
            GateSetPredicate(self._GATE_SET),
            DefaultRegisterPredicate(),
        ]

    def rebase_pass(self) -> BasePass:
        return auto_rebase_pass({OpType.CZ, OpType.Rz, OpType.Rx})

    def default_compilation_pass(self, optimisation_level: int = 1) -> BasePass:
        assert optimisation_level in range(3)
        passlist = [DecomposeBoxes(), FlattenRegisters()]
        if optimisation_level == 1:
            passlist.append(SynthesiseTket())
        elif optimisation_level == 2:
            passlist.append(FullPeepholeOptimise())
        passlist.append(self.rebase_pass())
        if optimisation_level > 0:
            passlist.append(EulerAngleReduction(OpType.Rx, OpType.Rz))
        return SequencePass(passlist)

    @property
    def _result_id_type(self) -> _ResultIdTuple:
        return (int,)

    def process_circuits(
        self,
        circuits: Iterable[Circuit],
        n_shots: Optional[Union[int, Sequence[int]]] = None,
        valid_check: bool = True,
        **kwargs: KwargTypes,
    ) -> List[ResultHandle]:
        handle_list = []
        if valid_check:
            self._check_all_circuits(circuits)
        for circuit in circuits:
            p = tk_to_pyquil(circuit)
            for qb in circuit.qubits:
                # Qubits with no gates will not be included in the Program
                # Add identities to ensure all qubits are present and dimension
                # is as expected
                p += I(Qubit_(qb.index[0]))
            handle = ResultHandle(uuid4().int)
            state = np.array(self._sim.wavefunction(p).amplitudes)
            try:
                phase = float(circuit.phase)
                coeff = np.exp(phase * np.pi * 1j)
                state *= coeff
            except ValueError:
                warning(
                    "Global phase is dependent on a symbolic parameter, so cannot "
                    "adjust for phase"
                )
            implicit_perm = circuit.implicit_qubit_permutation()
            res_qubits = [
                implicit_perm[qb] for qb in sorted(circuit.qubits, reverse=True)
            ]
            res = BackendResult(q_bits=res_qubits, state=state)
            self._cache[handle] = {"result": res}
            handle_list.append(handle)
        return handle_list

    def circuit_status(self, handle: ResultHandle) -> CircuitStatus:
        if handle in self._cache:
            return CircuitStatus(StatusEnum.COMPLETED)
        raise CircuitNotRunError(handle)

    def _gen_PauliTerm(self, term: QubitPauliString, coeff: complex = 1.0) -> PauliTerm:
        pauli_term = ID() * coeff
        for q, p in term.map.items():
            pauli_term *= PauliTerm(p.name, _default_q_index(q))
        return pauli_term  # type: ignore

    def get_pauli_expectation_value(
        self, state_circuit: Circuit, pauli: QubitPauliString
    ) -> complex:
        """Calculates the expectation value of the given circuit using the built-in QVM
        functionality

        :param state_circuit: Circuit that generates the desired state
            :math:`\\left|\\psi\\right>`.
        :type state_circuit: Circuit
        :param pauli: Pauli operator
        :type pauli: QubitPauliString
        :return: :math:`\\left<\\psi | P | \\psi \\right>`
        :rtype: complex
        """
        prog = tk_to_pyquil(state_circuit)
        pauli_term = self._gen_PauliTerm(pauli)
        return complex(self._sim.expectation(prog, [pauli_term]))

    def get_operator_expectation_value(
        self, state_circuit: Circuit, operator: QubitPauliOperator
    ) -> complex:
        """Calculates the expectation value of the given circuit with respect to the
        operator using the built-in QVM functionality

        :param state_circuit: Circuit that generates the desired state
            :math:`\\left|\\psi\\right>`.
        :type state_circuit: Circuit
        :param operator: Operator :math:`H`.
        :type operator: QubitPauliOperator
        :return: :math:`\\left<\\psi | H | \\psi \\right>`
        :rtype: complex
        """
        prog = tk_to_pyquil(state_circuit)
        pauli_sum = PauliSum(
            [self._gen_PauliTerm(term, coeff) for term, coeff in operator._dict.items()]
        )
        return complex(self._sim.expectation(prog, pauli_sum))

_xcirc = Circuit(1).Rx(1, 0)
_xcirc.add_phase(0.5)