"""
Contains the Structure class. This is used to represent an arbitrary
3-dimensional chemical species, which may be a complete molecule or
a fragment.

3D coordinates can be generated, and atoms can be fixed.
"""

from typing import List, Tuple, Dict
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Geometry import Point3D
import numpy as np

from pygen_structures.mol_containers.atom import Atom
from pygen_structures._functions_const import (
    bonds_to_adjacency,
    BOND_ORDER_TO_TYPE,
    get_expected_valence
)

class Structure:
    """
    An arbitrary chemical structure. This may be a complete molecule,
    but could also be an isolated residue fragment.
    """
    def __init__(
            self,
            atoms: List[Atom],
            bonds: List[Tuple[int, int]],
            improper_ics: (Dict[Tuple[int, ...], float], None) = None,
            fixed_atoms: (List[int], None) = None,
            set_charges: bool = True,
            use_etkdg: bool = True,
    ):
        self.atoms = atoms
        self.adjacency_matrix = bonds_to_adjacency(bonds)
        self.bonds = bonds
        if not fixed_atoms:
            fixed_atoms = list()
        self.fixed_atoms = fixed_atoms
        if not improper_ics:
            improper_ics = dict()
        self.improper_ics = improper_ics
        self.mol: (None, Chem.RWMol) = None
        self.use_etkdg = use_etkdg

        self.set_formal_charges()
        self.to_mol()
        self.reorder_atoms()
        self.generate_coordinates()
        if set_charges is False:
            for rd_atom in self.mol.GetAtoms():
                rd_atom.SetFormalCharge(0)

    def set_formal_charges(self) -> None:
        """
        Set Atom formal charges from valence and bond order.
        """
        for atom, connections in zip(self.atoms, self.adjacency_matrix):
            total_valence = np.sum(connections).item()
            expected_valence = get_expected_valence(atom.element_symbol)
            formal_charge = total_valence - expected_valence
            if formal_charge:
                atom.formal_charge = formal_charge

    def to_mol(self) -> Chem.RWMol:
        """
        Create an RDKit mol from the atoms and adjacency matrix.
        All stereocenters are set to clockwise initially.
        """
        if self.mol:
            return self.mol

        # noinspection PyArgumentList
        mol = Chem.RWMol()
        for atom in self.atoms:
            residue_index = atom.residue_number - 1
            rd_atom = Chem.Atom(atom.element_symbol)
            rd_atom.SetProp("atom_name", atom.atom_name)
            rd_atom.SetProp("atom_type", atom.atom_type)
            rd_atom.SetDoubleProp("partial_charge", atom.partial_charge)
            rd_atom.SetProp("residue_name", atom.residue_name)
            rd_atom.SetIntProp("residue_index", residue_index)
            atom_id = "{}:{}".format(residue_index, atom.atom_name)
            rd_atom.SetProp("atom_id", atom_id)
            rd_atom.SetFormalCharge(atom.formal_charge)
            mol.AddAtom(rd_atom)

        for bond in np.argwhere(self.adjacency_matrix):
            idx_x, idx_y = bond.tolist()
            if idx_y < idx_x:
                continue
            bond_order = self.adjacency_matrix[idx_x, idx_y]
            bond_type = BOND_ORDER_TO_TYPE[bond_order]
            mol.AddBond(idx_x, idx_y, bond_type)

        Chem.SanitizeMol(mol)
        chiral_centers = Chem.FindMolChiralCenters(mol, True, True)
        chiral_centers = [x[0] for x in chiral_centers]
        for atom_index in chiral_centers:
            rd_atom = mol.GetAtomWithIdx(atom_index)
            rd_atom.SetChiralTag(Chem.CHI_TETRAHEDRAL_CW)

        self.mol = mol
        return mol

    def generate_coordinates(self) -> None:
        """
        Generate a 3D conformation for the Structure using RDKit.
        This may take a few attempts to set stereochemistry correctly.
        """
        mol = self.to_mol()

        fixed_coords = {}
        for atom_index in self.fixed_atoms:
            atom = self.atoms[atom_index]
            coordinates = (atom.x, atom.y, atom.z)
            fixed_coords[atom_index] = Point3D(*coordinates)
        if self.use_etkdg:
            embed_parameters = AllChem.ETKDGv2()
        else:
            # noinspection PyArgumentList
            embed_parameters = AllChem.EmbedParameters()
            embed_parameters.useBasicKnowledge = True
            embed_parameters.ignoreSmoothingFailures = True
            embed_parameters.verbose = True
        if fixed_coords:
            embed_parameters.coordMap = fixed_coords
        AllChem.EmbedMolecule(mol, embed_parameters)

        centers_to_invert, previous_centers = set(), set()
        chiral_centers = Chem.FindMolChiralCenters(mol, True, True)
        chiral_centers = [x[0] for x in chiral_centers]

        n_cycles = 0
        while n_cycles < 100:
            n_cycles += 1
            # noinspection PyArgumentList
            conformer = mol.GetConformer()
            for (i, j, k, l), improper in self.improper_ics.items():
                if k not in chiral_centers:
                    continue
                measured_impr = AllChem.GetDihedralDeg(conformer, i, j, k, l)
                if (
                        (improper > 0 and measured_impr < 0) or
                        (improper < 0 and measured_impr > 0)
                ):
                    centers_to_invert.add(k)
            if not centers_to_invert:
                break
            # In case the chirality of one center is dependent upon another
            # This doesn't happen very often.
            if centers_to_invert == previous_centers:
                # Randomly pop one center and randomly add neighboring
                # chiral centers to the mix.
                center = centers_to_invert.pop()
                rd_atom = mol.GetAtomWithIdx(center)
                for neighbor in rd_atom.GetNeighbors():
                    neighbor_index = neighbor.GetIdx()
                    if neighbor_index in chiral_centers:
                        if np.random.random() >= 0.5:
                            centers_to_invert.add(neighbor_index)
                if not centers_to_invert:
                    centers_to_invert.add(center)
            for atom_index in centers_to_invert:
                rd_atom = mol.GetAtomWithIdx(atom_index)
                rd_atom.InvertChirality()
            previous_centers, centers_to_invert = centers_to_invert, set()
            AllChem.EmbedMolecule(mol, embed_parameters)
        else:
            raise ValueError("Unable to set stereochemistry correctly.")

        # noinspection PyArgumentList
        conformer = mol.GetConformer()
        for atom_index, atom in enumerate(self.atoms):
            position = conformer.GetAtomPosition(atom_index)
            atom.x = position.x
            atom.y = position.y
            atom.z = position.z

    def reorder_atoms(self) -> None:
        """
        Reorder the atoms and adjacency matrix to put the hydrogen
        atoms after their parent heavy atom. This is probably an easier
        layout for a person to follow, should they look at the file.
        """
        new_atom_order: List[int] = []
        for rd_atom in self.mol.GetAtoms():
            if rd_atom.GetSymbol() == "H":
                continue
            index = rd_atom.GetIdx()
            new_atom_order.append(index)
            neighbors = rd_atom.GetNeighbors()
            hydrogens = [n for n in neighbors if n.GetSymbol() == "H"]
            hydrogen_indices = [n.GetIdx() for n in hydrogens]
            new_atom_order.extend(hydrogen_indices)
        self.mol = Chem.RenumberAtoms(self.mol, new_atom_order)

        self.atoms[:] = [self.atoms[x] for x in new_atom_order]
        for new_index, atom in enumerate(self.atoms):
            atom.atom_serial = new_index + 1

        for row in self.adjacency_matrix:
            row[:] = [row[x] for x in new_atom_order]
        for column in self.adjacency_matrix.T:
            column[:] = [column[x] for x in new_atom_order]

        for old_improper_indices in list(self.improper_ics.keys()):
            (i, j, k, l) = old_improper_indices
            i, j = new_atom_order[i], new_atom_order[j]
            k, l = new_atom_order[k], new_atom_order[l]
            improper = self.improper_ics.pop(old_improper_indices)
            self.improper_ics[(i, j, k, l)] = improper


# -*- coding: utf-8 -*-

import os
import requests
import glob
import numpy as np
import pickle
import torch as th
import dgl
import subprocess
from dgl.nn.pytorch import GraphConv
import torch.nn as nn
import torch.nn.functional as F
from pathlib import Path
from passer import GCNN

class Passer(object):
    def __init__(self, pdbID, pdbFile, chain, save):
        self.PDBFILE = pdbFile
        # indicate whether user have input PDB file
        self.UPLOAD = False
        self.save = True if save == "Y" or save == "y" else False
        self.PDB = pdbID
        self.pdbDIRECTION = "./"
        if self.PDBFILE:
            self.UPLOAD = True
            slashIndex = self.PDBFILE.rfind("/")
            self.PDB = self.PDBFILE[slashIndex + 1:].split(".")[0]
            self.pdbDIRECTION = self.PDBFILE[:slashIndex + 1]
        self.pocketDirection = "%s%s_out/pockets/"\
            %(self.pdbDIRECTION, self.PDB)
        self.CHAIN = chain
        self.HOME = str(Path.home()) + "/"
        self.modelDirection = self.HOME + ".passerModels/"

    def __download(self):
        self.URL = "https://files.rcsb.org/download/" + self.PDB + ".pdb"
        if self.PDB:
            r = requests.get(url = self.URL)
            open("%s.pdb" %self.PDB, "wb").write(r.content)

    def __runFPocket(self):
        if self.CHAIN:
            subprocess.call(["fpocket", "-f", "%s%s.pdb" \
                %(self.pdbDIRECTION, self.PDB), "-k", "%s" %self.CHAIN])
        else:
            subprocess.call(["fpocket", "-f", "%s%s.pdb" \
                %(self.pdbDIRECTION, self.PDB)])
        self.fileDirection = "%s%s_out/%s_info.txt" \
            %(self.pdbDIRECTION, self.PDB, self.PDB)

    def __extractPocket(self):
        pocket = open(self.fileDirection + "", "r").readlines()
        pocket_num = len(pocket) // 21
        features = []
        for index in range(pocket_num):
            cur_feature = []
            cur = pocket[index * 21 : (index + 1) * 21]
            for line in cur[1:-1]:
                cur_feature.append(float(line.split("\t")[2][:-1]))
            if len(cur_feature) == 18:
                warnings.warn("potential error in features of %s" \
                    %self.fileDirection)
            features.append(cur_feature)
        return features

    def __rank(self, cur_feature, index):
        score = [cur_feature[m][index] for m in range(len(cur_feature))]
        score = sorted(score, reverse = True)
        for p in range(len(cur_feature)):
            cur_feature[p].append(score.index(cur_feature[p][index]) + 1)
        return cur_feature

    def __collectFPocket(self):
        self.features = self.__extractPocket()
        # add additional ranking features
        self.features = np.array(self.features)

    def __graph(self, fileDirection, BOND_THRESHOLD = 10):
        # define atom and bond value
        atom_map = {"C": 12, "N": 14, "O": 16, "S": 32}
        pocket = open(fileDirection, "r").readlines()
        coordinates = []
        atom_type = []

        for line in pocket:
            if line[:4] == "ATOM":
                info = line.split()
                if len(info[-2]) == 1:
                    atom_type.append(info[-2])
                else:
                    atom_type.append(info[-1])
                try:
                    coordinates.append(list(map(float, info[6:9])))
                except:
                    try:
                        dot_index = info[7].index(".")
                        coordinates.append([float(info[6]), \
                            float(info[7][:dot_index+4]), \
                            float(info[7][dot_index+4:])])
                    except:
                        dot_index = info[6].index(".")
                        coordinates.append([float(info[6][:dot_index+4]), \
                            float(info[6][dot_index+4:]), float(info[7])])

        coordinates = np.array(coordinates)

        NUM_OF_NODES = len(atom_type)
        start_node = []
        end_node = []

        for i in range(NUM_OF_NODES):
            connected_index = []
            for j in range(i + 1, NUM_OF_NODES):
                distance = np.sqrt(sum(np.square(coordinates[i] \
                    - coordinates[j])))
                # if distance is within threshold
                if distance <= BOND_THRESHOLD:
                    connected_index.append(j)
            if connected_index:
                start_node.extend([i] * len(connected_index))
                end_node.extend(connected_index)

        u = start_node + end_node
        v = end_node + start_node

        g = dgl.DGLGraph()
        g.add_nodes(NUM_OF_NODES)
        g.add_edges(u, v)
        return g

    def __collectPockets(self):
        fileNames = glob.glob(self.pocketDirection + "*.pdb")
        fileNames = sorted(fileNames, key = \
            lambda x : int(x.split("pocket")[-1].split("_")[0]))
        self.pockets = []
        for fileN in fileNames:
            self.pockets.append(self.__graph(fileN, 10))

    def __collectModels(self):
        self.xgbmodels = pickle.load(open(self.modelDirection + "xgboost.pkl", "rb"))

        ## collect GCNN models
        model = GCNN.Classifier(1, 256, 2)
        self.gcnnmodels = []
        checkpoint = th.load(self.modelDirection + "GCNN.pt")
        for param in checkpoint.values():
            model.load_state_dict(param)
            self.gcnnmodels.append(model)

    def __extractResidue(self):
        pocketFiles = glob.glob(self.pocketDirection + "*.pdb")
        pocketFiles = sorted(pocketFiles, key = lambda x : \
            int(x[x.rfind("pocket") + 6:x.rfind("_")]))
        residues = []
        for index in range(min(3, len(pocketFiles))):
            curFile = open(pocketFiles[index], "r").readlines()
            curResidues = []
            for line in curFile:
                if line[:4] == "ATOM":
                    info = line.split()
                    curResidues.append(info[4] + ":" + info[3] + info[5])
            residues.append(" ".join(set(curResidues)))
        return residues

    def predict(self):
        self.__download()
        self.__runFPocket()
        self.__collectFPocket()
        self.__collectPockets()
        self.__collectModels()

        # predict using each model
        self.xgboostProbs = []
        for model in self.xgbmodels:
            self.xgboostProbs.append(model.predict_proba(self.features))

        self.xgboostProbs = np.array(self.xgboostProbs)
        self.xgboostProbs = np.mean(self.xgboostProbs, axis = 0)

        self.gcnnProbs = []
        for model in self.gcnnmodels:
            # Convert a list of tuples to two lists
            test_bg = dgl.batch(self.pockets)
            probs_Y = th.softmax(model(test_bg), 1)
            self.gcnnProbs.append(probs_Y.tolist())

        self.gcnnProbs = np.array(self.gcnnProbs)
        self.gcnnProbs = np.mean(self.gcnnProbs, axis = 0)

        # combine probabilities
        self.meanValues = (self.xgboostProbs + self.gcnnProbs) / 2
        residues = self.__extractResidue()

        for i in range(0, min(len(self.meanValues), 3)):
            print("Pocket %d: %.1f%%" %(i+1, self.meanValues[i][1]*100))
            print(residues[i])

        # remove intermediate files
        if not self.UPLOAD:
            os.system("rm %s.pdb" %self.PDB)
        if not self.save:
            os.system("rm -r %s%s_out" %(self.pdbDIRECTION, self.PDB))


import logging
from dataclasses import fields

import numpy as np
import pandas as pd
import pyspark

from sm.engine.annotation.formula_validator import Metrics
from sm.engine.annotation.png_generator import PngGenerator
from sm.engine.config import SMConfig
from sm.engine.db import DB
from sm.engine.image_storage import ImageStorage
from sm.engine.ion_mapping import get_ion_id_mapping
from sm.engine.utils.numpy_json_encoder import numpy_json_dumps

logger = logging.getLogger('engine')
METRICS_INS = (
    'INSERT INTO annotation ('
    '   job_id, formula, chem_mod, neutral_loss, adduct, msm, fdr, stats, iso_image_ids, ion_id'
    ') '
    'VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)'
)
METRICS_FIELDS = [
    *[f.name for f in [*fields(Metrics)] if f.name != 'formula_i' and f.name != 'msm'],
    'chaos_fdr',
    'spatial_fdr',
    'spectral_fdr',
    'mz_err_abs_fdr',
    'mz_err_rel_fdr',
]

class SearchResults:
    """Container for molecule search results."""

    def __init__(self, ds_id: str, job_id: int, n_peaks: int, charge: int):
        """
        Args:
            ds_id: dataset id
            job_id: annotation job id
            n_peaks: number of isotopic peaks
            charge: charge of ions
        """
        self.ds_id = ds_id
        self.job_id = job_id
        self.n_peaks = n_peaks
        self.charge = charge

    @staticmethod
    def _metrics_table_row_gen(job_id, metr_df, ion_image_ids, ion_mapping):
        stats_cols = [f for f in METRICS_FIELDS if f in metr_df.columns]
        for _, row in metr_df.iterrows():
            metr_json = numpy_json_dumps({m: row[m] for m in stats_cols})
            if row.formula_i not in ion_image_ids:
                logger.debug(f'Missing "formula_i": {row}, {ion_image_ids}')
            image_ids = ion_image_ids[row.formula_i]
            yield (
                job_id,
                row.formula,
                row.chem_mod,
                row.neutral_loss,
                row.adduct,
                float(row.msm),
                float(row.fdr),
                metr_json,
                image_ids,
                ion_mapping[row.formula, row.chem_mod, row.neutral_loss, row.adduct],
            )

    def store_ion_metrics(self, ion_metrics_df, ion_image_ids, db):
        """Store ion metrics and iso image ids in the database."""

        logger.info('Storing iso image metrics')

        ions = ion_metrics_df[['formula', 'chem_mod', 'neutral_loss', 'adduct']]
        ion_tuples = list(ions.itertuples(False, None))
        ion_mapping = get_ion_id_mapping(db, ion_tuples, self.charge)

        rows = self._metrics_table_row_gen(
            self.job_id, ion_metrics_df.reset_index(), ion_image_ids, ion_mapping
        )
        db.insert(METRICS_INS, list(rows))

    def _post_images_to_image_store(self, ion_images_rdd, alpha_channel, n_peaks):
        logger.info('Posting iso images to image store')
        png_generator = PngGenerator(alpha_channel, greyscale=True)
        ds_id = self.ds_id
        sm_config = SMConfig.get_conf()

        def generate_png_and_post(partition):
            image_storage = ImageStorage(sm_config)

            for formula_i, imgs in partition:
                iso_image_ids = [None] * n_peaks
                for k, img in enumerate(imgs):
                    if img is not None:
                        img_bytes = png_generator.generate_png(img.toarray())
                        iso_image_ids[k] = image_storage.post_image(
                            image_storage.ISO, ds_id, img_bytes
                        )

                yield formula_i, iso_image_ids

        return dict(ion_images_rdd.mapPartitions(generate_png_and_post).collect())

    def store(
        self,
        metrics_df: pd.DataFrame,
        formula_images_rdd: pyspark.RDD,
        alpha_channel: np.ndarray,
        db: DB,
    ):
        """Store ion metrics and iso images.

        Args:
            metrics_df: formula, adduct, msm, fdr, individual metrics
            formula_images_rdd: collection of 2d intensity arrays (in coo_matrix format)
            alpha_channel: Image alpha channel (2D, 0..1)
            db: database connection
        """
        logger.info('Storing search results to the DB')
        ion_image_ids = self._post_images_to_image_store(
            formula_images_rdd, alpha_channel, self.n_peaks
        )
        self.store_ion_metrics(metrics_df, ion_image_ids, db)


from __future__ import division, absolute_import, print_function
from past.builtins import xrange

import numpy as np
import os
import sys
import esutil

import matplotlib.pyplot as plt

from .fgcmUtilities import zpFlagDict
from .fgcmUtilities import expFlagDict

from .sharedNumpyMemManager import SharedNumpyMemManager as snmm

## FIXME with plotter!
## also note retrieval!

class FgcmZeropoints(object):
    """
    Class to compute final zeropoints

    parameters
    ----------
    fgcmConfig: FgcmConfig
    fgcmPars: FgcmParameters
    fgcmLUT: FgcmLUT
    fgcmGray: FgcmGray
    fgcmRetrieval: FgcmRetrieval

    Config variables
    ----------------
    minCCDPerExp: int
       Minimum number of CCDs on an exposure to recover bad ccds
    minStarPerCCD: int
       Minimum number of stars to recover zeropoint through matching
    maxCCDGrayErr: float
       Maximum CCD Gray error to recover zeropoint through matching
    sigma0Cal: float
       Minimum error on zeropoint allowed
    expGrayRecoverCut: float
       Minimum (negative!) exposure gray to consider recovering via focal-plane average
    expGrayErrRecoverCut: float
       Maximum exposure gray error to consider recovering via focal-plane average
    expVarGrayPhotometricCut: float
       Exposure gray variance to consider recovering via focal-plane average
    """

    def __init__(self,fgcmConfig,fgcmPars,fgcmLUT,fgcmGray,fgcmRetrieval):

        self.fgcmPars = fgcmPars
        self.fgcmLUT = fgcmLUT
        self.fgcmGray = fgcmGray
        self.fgcmRetrieval = fgcmRetrieval

        self.fgcmLog = fgcmConfig.fgcmLog

        self.fgcmLog.info('Initializing FgcmZeropoints...')

        self.illegalValue = fgcmConfig.illegalValue
        self.outputPath = fgcmConfig.outputPath
        self.cycleNumber = fgcmConfig.cycleNumber
        self.outfileBaseWithCycle = fgcmConfig.outfileBaseWithCycle
        self.plotPath = fgcmConfig.plotPath
        self.zptAB = fgcmConfig.zptAB
        self.ccdStartIndex = fgcmConfig.ccdStartIndex
        self.ccdOffsets = fgcmConfig.ccdOffsets
        self.minCCDPerExp = fgcmConfig.minCCDPerExp
        self.minStarPerCCD = fgcmConfig.minStarPerCCD
        self.maxCCDGrayErr = fgcmConfig.maxCCDGrayErr
        self.sigma0Cal = fgcmConfig.sigma0Cal
        self.expField = fgcmConfig.expField
        self.ccdField = fgcmConfig.ccdField
        self.expGrayRecoverCut = fgcmConfig.expGrayRecoverCut
        self.expGrayErrRecoverCut = fgcmConfig.expGrayErrRecoverCut
        self.expVarGrayPhotometricCut = fgcmConfig.expVarGrayPhotometricCut

    def computeZeropoints(self):
        """
        Compute the zeropoints from all the fits that have been performed.

        Output attributes
        -----------------
        zpStruct: Zero point recarray (nExp * nCCD)
           expField: Exposure field name
           ccdField: CCD field name
           'FGCM_FLAG': Quality flag value
           'FGCM_ZPT': Zeropoint
           'FGCM_ZPTERR': Error on zeropoint
           'FGCM_I0': I0 for exp/ccd (throughput)
           'FGCM_I10': I10 for exp/ccd (chromatic)
           'FGCM_R0': Retrieved throughput integral
           'FGCM_R10': Retrieved chromatic integral
           'FGCM_GRY': Delta-zeropoint due to CCD gray
           'FGCM_ZPTVAR': Variance of zeropoint as estimated from gray corrections
           'FGCM_TILINGS': Average number of observations of calib stars on CCD
           'FGCM_FPGRY': Average focal-plane gray (exp_gray)
           'FGCM_FPVAR': Focal-plane variance
           'FGCM_DUST': Delta-zeropoint due to mirror/corrector dust buildup
           'FGCM_FLAT': Delta-zeropoint due to superStarFlat
           'FGCM_APERCORR': Delta-zeropoint due to aperture correction
           'EXPTIME': Exposure time (seconds)
           'FILTERNAME': Filter name
           'BAND': band name
        atmStruct: Atmosphere parameter recarray (nExp)
           'PMB': Barometric pressure (mb)
           'PWV': preciptable water vapor (mm)
           'TAU': aerosol optical index
           'ALPHA': Aerosol slope
           'O3': Ozone (dob)
           'SECZENITH': secant(zenith angle)
        """

        # first, we need to get relevant quantities from shared memory.
        expGray = snmm.getArray(self.fgcmGray.expGrayHandle)
        expGrayErr = snmm.getArray(self.fgcmGray.expGrayErrHandle)
        expGrayRMS = snmm.getArray(self.fgcmGray.expGrayRMSHandle)
        expNGoodCCDs = snmm.getArray(self.fgcmGray.expNGoodCCDsHandle)
        expNGoodTilings = snmm.getArray(self.fgcmGray.expNGoodTilingsHandle)
        expGrayColorSplit = snmm.getArray(self.fgcmGray.expGrayColorSplitHandle)
        expGrayRMSColorSplit = snmm.getArray(self.fgcmGray.expGrayRMSColorSplitHandle)

        ccdGray = snmm.getArray(self.fgcmGray.ccdGrayHandle)
        ccdGrayRMS = snmm.getArray(self.fgcmGray.ccdGrayRMSHandle)
        ccdGrayErr = snmm.getArray(self.fgcmGray.ccdGrayErrHandle)
        ccdNGoodStars = snmm.getArray(self.fgcmGray.ccdNGoodStarsHandle)
        ccdNGoodTilings = snmm.getArray(self.fgcmGray.ccdNGoodTilingsHandle)

        r0 = snmm.getArray(self.fgcmRetrieval.r0Handle)
        r10 = snmm.getArray(self.fgcmRetrieval.r10Handle)

        # and we need to make sure we have the parameters, and
        #  set these to the exposures
        parArray = self.fgcmPars.getParArray(fitterUnits=False)
        # Use *input* values for the retrieved atmosphere parameters (if needed)
        self.fgcmPars.parsToExposures(retrievedInput=True)

        # set up output structures

        self.fgcmLog.info('Building zeropoint structure...')

        zpStruct = np.zeros(self.fgcmPars.nExp*self.fgcmPars.nCCD,
                            dtype=[(self.expField,'i4'), # done
                                   (self.ccdField,'i2'), # done
                                   ('FGCM_FLAG','i2'), # done
                                   ('FGCM_ZPT','f8'), # done
                                   ('FGCM_ZPTERR','f8'), # done
                                   ('FGCM_I0','f8'), # done
                                   ('FGCM_I10','f8'), # done
                                   ('FGCM_R0','f8'),
                                   ('FGCM_R10','f8'),
                                   ('FGCM_GRY','f8'), # done
                                   ('FGCM_ZPTVAR','f8'), # done
                                   ('FGCM_TILINGS','f8'), # done
                                   ('FGCM_FPGRY','f8'), # done
                                   ('FGCM_FPVAR','f8'), # done
                                   ('FGCM_FPGRY_CSPLIT', 'f8', 3),
                                   ('FGCM_FPGRY_CSPLITVAR', 'f8', 3),
                                   ('FGCM_DUST','f8'), # done
                                   ('FGCM_FLAT','f8'), # done
                                   ('FGCM_APERCORR','f8'), # done
                                   ('EXPTIME','f4'), # done
                                   ('FILTERNAME','a2'),
                                   ('BAND','a2')]) # done

        atmStruct = np.zeros(self.fgcmPars.nExp,
                             dtype=[(self.expField,'i4'),
                                    ('PMB','f8'),
                                    ('PWV','f8'),
                                    ('TAU','f8'),
                                    ('ALPHA','f8'),
                                    ('O3','f8'),
                                    ('SECZENITH','f8')])

        ## start with zpStruct

        # fill out exposures and ccds
        zpStruct[self.expField][:] = np.repeat(self.fgcmPars.expArray,
                                          self.fgcmPars.nCCD)
        zpStruct[self.ccdField][:] = np.tile(np.arange(self.fgcmPars.nCCD)+self.ccdStartIndex,
                                        self.fgcmPars.nExp)

        # get the exposure indices and CCD indices
        zpExpIndex = np.searchsorted(self.fgcmPars.expArray,zpStruct[self.expField])
        zpCCDIndex = zpStruct[self.ccdField] - self.ccdStartIndex

        # fill exposure quantities
        lutFilterNameArray = np.array(self.fgcmPars.lutFilterNames)
        zpStruct['FILTERNAME'][:] = lutFilterNameArray[self.fgcmPars.expLUTFilterIndex[zpExpIndex]]
        bandArray = np.array(self.fgcmPars.bands)
        zpStruct['BAND'][:] = bandArray[self.fgcmPars.expBandIndex[zpExpIndex]]
        zpStruct['EXPTIME'][:] = self.fgcmPars.expExptime[zpExpIndex]

        # fill in the superstar flat
        zpStruct['FGCM_FLAT'][:] = self.fgcmPars.expCCDSuperStar[zpExpIndex,
                                                                 zpCCDIndex]

        # fill in the optics dust
        zpStruct['FGCM_DUST'][:] = self.fgcmPars.expQESys[zpExpIndex]

        # fill in the aperture correction
        zpStruct['FGCM_APERCORR'][:] = self.fgcmPars.expApertureCorrection[zpExpIndex]

        # fill in the retrieved values
        zpStruct['FGCM_R0'][:] = r0[zpExpIndex, zpCCDIndex]
        zpStruct['FGCM_R10'][:] = r10[zpExpIndex, zpCCDIndex]

        # and the focal-plane gray and var...
        # these are only filled in for those exposures where we have it computed
        zpStruct['FGCM_FPGRY'][:] = self.illegalValue
        zpStruct['FGCM_FPVAR'][:] = self.illegalValue

        zpExpOk, = np.where(expNGoodCCDs[zpExpIndex] >= self.minCCDPerExp)
        zpStruct['FGCM_FPGRY'][zpExpOk] = expGray[zpExpIndex[zpExpOk]]
        zpStruct['FGCM_FPVAR'][zpExpOk] = expGrayRMS[zpExpIndex[zpExpOk]]**2.

        zpStruct['FGCM_FPGRY_CSPLIT'][zpExpOk, :] = expGrayColorSplit[zpExpIndex[zpExpOk], :]
        zpStruct['FGCM_FPGRY_CSPLITVAR'][zpExpOk, :] = expGrayRMSColorSplit[zpExpIndex[zpExpOk], :]**2.

        self.fgcmLog.info('%d exposure/ccd sets have exposures with >=%d good ccds' %
                         (zpExpOk.size, self.minCCDPerExp))

        # look up the I0 and I10s.  These are defined for everything
        #  (even if only standard bandpass, it'll grab instrumental)

        ## FIXME: will probably need some sort of rotation information at some point

        ## FIXME: check that the signs are correct!
        # need secZenith for each exp/ccd pair
        ccdHA = (self.fgcmPars.expTelHA[zpExpIndex] -
                 np.radians(self.ccdOffsets['DELTA_RA'][zpCCDIndex]))
        ccdDec = (self.fgcmPars.expTelDec[zpExpIndex] +
                  np.radians(self.ccdOffsets['DELTA_DEC'][zpCCDIndex]))
        ccdSecZenith = 1./(np.sin(ccdDec) * self.fgcmPars.sinLatitude +
                           np.cos(ccdDec) * self.fgcmPars.cosLatitude * np.cos(ccdHA))

        # and do the LUT lookups
        lutIndices = self.fgcmLUT.getIndices(self.fgcmPars.expLUTFilterIndex[zpExpIndex],
                                             self.fgcmPars.expPWV[zpExpIndex],
                                             self.fgcmPars.expO3[zpExpIndex],
                                             #np.log(self.fgcmPars.expTau[zpExpIndex]),
                                             self.fgcmPars.expLnTau[zpExpIndex],
                                             self.fgcmPars.expAlpha[zpExpIndex],
                                             ccdSecZenith,
                                             zpCCDIndex,
                                             self.fgcmPars.expPmb[zpExpIndex])
        zpStruct['FGCM_I0'][:] = self.fgcmLUT.computeI0(self.fgcmPars.expPWV[zpExpIndex],
                                                        self.fgcmPars.expO3[zpExpIndex],
                                                        #np.log(self.fgcmPars.expTau[zpExpIndex]),
                                                        self.fgcmPars.expLnTau[zpExpIndex],
                                                        self.fgcmPars.expAlpha[zpExpIndex],
                                                        ccdSecZenith,
                                                        self.fgcmPars.expPmb[zpExpIndex],
                                                        lutIndices)
        zpStruct['FGCM_I10'][:] = self.fgcmLUT.computeI1(self.fgcmPars.expPWV[zpExpIndex],
                                                         self.fgcmPars.expO3[zpExpIndex],
                                                         #np.log(self.fgcmPars.expTau[zpExpIndex]),
                                                         self.fgcmPars.expLnTau[zpExpIndex],
                                                         self.fgcmPars.expAlpha[zpExpIndex],
                                                         ccdSecZenith,
                                                         self.fgcmPars.expPmb[zpExpIndex],
                                                         lutIndices) / zpStruct['FGCM_I0'][:]

        # Set the tilings, gray values, and zptvar

        zpStruct['FGCM_TILINGS'][:] = self.illegalValue
        zpStruct['FGCM_GRY'][:] = self.illegalValue
        zpStruct['FGCM_ZPTVAR'][:] = self.illegalValue

        goodCCD, = np.where((ccdNGoodStars[zpExpIndex,zpCCDIndex] >=
                             self.minStarPerCCD) &
                            (ccdGrayErr[zpExpIndex,zpCCDIndex] <=
                             self.maxCCDGrayErr))
        zpStruct['FGCM_TILINGS'][goodCCD] = ccdNGoodTilings[zpExpIndex[goodCCD],
                                                            zpCCDIndex[goodCCD]]
        zpStruct['FGCM_GRY'][goodCCD] = ccdGray[zpExpIndex[goodCCD],
                                                zpCCDIndex[goodCCD]]
        zpStruct['FGCM_ZPTVAR'][goodCCD] = ccdGrayErr[zpExpIndex[goodCCD],
                                                      zpCCDIndex[goodCCD]]**2.

        self.fgcmLog.info('%d CCDs are Good (>=%d stars; err <= %.3f)' %
                         (goodCCD.size, self.minStarPerCCD, self.maxCCDGrayErr))

        # check: if this has too few stars on the ccd OR the ccd error is too big
        #        AND the exposure has enough ccds
        #        AND the exposure gray error is small enough
        #        AND the exposure gray rms is small enough
        #        AND the exposure gray is not very large (configurable)
        #  then we can use the exposure stats to fill the variables

        badCCDGoodExp, = np.where(((ccdNGoodStars[zpExpIndex,zpCCDIndex] <
                            self.minStarPerCCD) |
                           (ccdGrayErr[zpExpIndex,zpCCDIndex] >
                            self.maxCCDGrayErr)) &
                          (expNGoodCCDs[zpExpIndex] >=
                           self.minCCDPerExp) &
                          (expGrayErr[zpExpIndex] <=
                           self.expGrayErrRecoverCut) &
                          (expGrayRMS[zpExpIndex] <=
                           np.sqrt(self.expVarGrayPhotometricCut)) &
                          (expGray[zpExpIndex] >=
                           self.expGrayRecoverCut))

        zpStruct['FGCM_TILINGS'][badCCDGoodExp] = expNGoodTilings[zpExpIndex[badCCDGoodExp]]
        zpStruct['FGCM_GRY'][badCCDGoodExp] = expGray[zpExpIndex[badCCDGoodExp]]
        zpStruct['FGCM_ZPTVAR'][badCCDGoodExp] = expGrayRMS[zpExpIndex[badCCDGoodExp]]**2.

        self.fgcmLog.info('%d CCDs recovered from good exposures (>=%d good CCDs, etc.)' %
                         (badCCDGoodExp.size, self.minCCDPerExp))

        # flag the photometric (fit) exposures
        photZpIndex, = np.where(self.fgcmPars.expFlag[zpExpIndex] == 0)

        photFitBand, = np.where(~self.fgcmPars.expExtraBandFlag[zpExpIndex[photZpIndex]])
        zpStruct['FGCM_FLAG'][photZpIndex[photFitBand]] |= (
            zpFlagDict['PHOTOMETRIC_FIT_EXPOSURE'])

        self.fgcmLog.info('%d CCDs marked as photometric, used in fit' %
                         (photFitBand.size))

        photExtraBand, = np.where(self.fgcmPars.expExtraBandFlag[zpExpIndex[photZpIndex]])
        zpStruct['FGCM_FLAG'][photZpIndex[photExtraBand]] |= (
            zpFlagDict['PHOTOMETRIC_EXTRA_EXPOSURE'])

        self.fgcmLog.info('%d CCDs marked as photometric, not used in fit' %
                         (photExtraBand.size))

        # flag the non-photometric exposures on calibratable nights
        rejectMask = (expFlagDict['TOO_FEW_EXP_ON_NIGHT'] |
                      expFlagDict['BAND_NOT_IN_LUT'] |
                      expFlagDict['NO_STARS'])
        acceptMask = (expFlagDict['EXP_GRAY_TOO_NEGATIVE'] |
                      expFlagDict['EXP_GRAY_TOO_POSITIVE'] |
                      expFlagDict['VAR_GRAY_TOO_LARGE'] |
                      expFlagDict['TOO_FEW_STARS'])

        nonPhotZpIndex, = np.where(((self.fgcmPars.expFlag[zpExpIndex] & acceptMask) > 0) &
                           ((self.fgcmPars.expFlag[zpExpIndex] & rejectMask) == 0))
        zpStruct['FGCM_FLAG'][nonPhotZpIndex] |= zpFlagDict['NONPHOTOMETRIC_FIT_NIGHT']

        self.fgcmLog.info('%d CCDs marked non-photometric, on a night with a fit' %
                         (nonPhotZpIndex.size))

        # and the exposures on non-calibratable nights (photometric or not, we don't know)
        rejectMask = (expFlagDict['NO_STARS'] |
                      expFlagDict['BAND_NOT_IN_LUT'])
        acceptMask = (expFlagDict['TOO_FEW_EXP_ON_NIGHT'])
        badNightZpIndex, = np.where(((self.fgcmPars.expFlag[zpExpIndex] & acceptMask) > 0) &
                                    ((self.fgcmPars.expFlag[zpExpIndex] & rejectMask) == 0))
        zpStruct['FGCM_FLAG'][badNightZpIndex] |= zpFlagDict['NOFIT_NIGHT']

        self.fgcmLog.info('%d CCDs on nights without a fit (assume standard atmosphere)' %
                         (badNightZpIndex.size))

        # and finally, the hopeless exposures
        acceptMask = (expFlagDict['NO_STARS'] |
                      expFlagDict['BAND_NOT_IN_LUT'])
        hopelessZpIndex, = np.where(((self.fgcmPars.expFlag[zpExpIndex] & acceptMask) > 0))
        zpStruct['FGCM_FLAG'][hopelessZpIndex] |= zpFlagDict['CANNOT_COMPUTE_ZEROPOINT']

        self.fgcmLog.info('%d CCDs marked as hopeless (cannot compute zeropoint)' %
                         (hopelessZpIndex.size))

        # now we can fill the zeropoints

        zpStruct['FGCM_ZPT'][:] = self.illegalValue
        zpStruct['FGCM_ZPTERR'][:] = self.illegalValue

        # start with the passable flag 1,2,4 exposures

        acceptMask = (zpFlagDict['PHOTOMETRIC_FIT_EXPOSURE'] |
              zpFlagDict['PHOTOMETRIC_EXTRA_EXPOSURE'] |
              zpFlagDict['NONPHOTOMETRIC_FIT_NIGHT'])

        okZpIndex, = np.where((zpStruct['FGCM_FLAG'] & acceptMask) > 0)

        okCCDZpIndexFlag = ((zpStruct['FGCM_I0'][okZpIndex] > 0.0) &
                    (zpStruct['FGCM_FLAT'][okZpIndex] > self.illegalValue) &
                    (zpStruct['FGCM_DUST'][okZpIndex] > self.illegalValue) &
                    (zpStruct['FGCM_APERCORR'][okZpIndex] > self.illegalValue) &
                    (zpStruct['FGCM_GRY'][okZpIndex] > self.illegalValue))

        okCCDZpIndex = okZpIndex[okCCDZpIndexFlag]

        self._computeZpt(zpStruct,okCCDZpIndex)
        self._computeZptErr(zpStruct,zpExpIndex,okCCDZpIndex)

        badCCDZpExp = okZpIndex[~okCCDZpIndexFlag]
        zpStruct['FGCM_FLAG'][badCCDZpExp] |=  zpFlagDict['TOO_FEW_STARS_ON_CCD']

        # and the flag 8 extra exposures

        acceptMask = zpFlagDict['NOFIT_NIGHT']

        mehZpIndex, = np.where((zpStruct['FGCM_FLAG'] & acceptMask) > 0)

        mehCCDZpIndexFlag = ((zpStruct['FGCM_I0'][mehZpIndex] > 0.0) &
                             (zpStruct['FGCM_FLAT'][mehZpIndex] > self.illegalValue) &
                             (zpStruct['FGCM_DUST'][mehZpIndex] > self.illegalValue) &
                             (zpStruct['FGCM_APERCORR'][mehZpIndex] > self.illegalValue) &
                             (zpStruct['FGCM_GRY'][mehZpIndex] > self.illegalValue) &
                             (ccdNGoodStars[zpExpIndex[mehZpIndex],zpCCDIndex[mehZpIndex]] >=
                              self.minStarPerCCD) &
                             (ccdGrayErr[zpExpIndex[mehZpIndex],zpCCDIndex[mehZpIndex]] <=
                              self.maxCCDGrayErr))

        mehCCDZpIndex = mehZpIndex[mehCCDZpIndexFlag]

        self._computeZpt(zpStruct,mehCCDZpIndex)
        self._computeZptErr(zpStruct,zpExpIndex,mehCCDZpIndex)

        badCCDZpExp = mehZpIndex[~mehCCDZpIndexFlag]
        zpStruct['FGCM_FLAG'][badCCDZpExp] |= zpFlagDict['TOO_FEW_STARS_ON_CCD']
        zpStruct['FGCM_FLAG'][badCCDZpExp] |= zpFlagDict['CANNOT_COMPUTE_ZEROPOINT']

        # record as a class element
        self.zpStruct = zpStruct

        #################################
        # and make the parameter file
        atmStruct[self.expField] = self.fgcmPars.expArray
        atmStruct['PMB'] = self.fgcmPars.expPmb
        atmStruct['PWV'] = self.fgcmPars.expPWV
        #atmStruct['TAU'] = self.fgcmPars.expTau
        atmStruct['TAU'] = np.exp(self.fgcmPars.expLnTau)
        atmStruct['ALPHA'] = self.fgcmPars.expAlpha
        atmStruct['O3'] = self.fgcmPars.expO3
        atmStruct['SECZENITH'] = 1./(np.sin(self.fgcmPars.expTelDec) *
                                     self.fgcmPars.sinLatitude +
                                     np.cos(self.fgcmPars.expTelDec) *
                                     self.fgcmPars.cosLatitude *
                                     np.cos(self.fgcmPars.expTelHA))

        # record as a class element
        self.atmStruct=atmStruct

        ############
        ## plots
        ############
        self.fgcmLog.info('Making I1/R1 plots...')

        # plotter = FgcmZeropointPlotter(zpStruct, self.fgcmPars.bands,
        #                               self.plotPath, self.outfileBaseWithCycle)
        plotter = FgcmZeropointPlotter(zpStruct, self.fgcmPars.lutFilterNames,
                                       self.plotPath, self.outfileBaseWithCycle)

        plotter.makeR1I1Plots()
        plotter.makeR1I1Maps(self.ccdOffsets, ccdField=self.ccdField)

        self.fgcmLog.info('Making zeropoint summary plots...')

        expZpMean = np.zeros(self.fgcmPars.nExp,dtype='f4')
        expZpNCCD = np.zeros(self.fgcmPars.nExp,dtype='i4')

        rejectMask = (zpFlagDict['CANNOT_COMPUTE_ZEROPOINT'] |
                      zpFlagDict['TOO_FEW_STARS_ON_CCD'])

        okCCD,=np.where((zpStruct['FGCM_FLAG'] & rejectMask) == 0)

        np.add.at(expZpMean,
                  zpExpIndex[okCCD],
                  zpStruct['FGCM_ZPT'][okCCD])
        np.add.at(expZpNCCD,
                  zpExpIndex[okCCD],
                  1)

        gd,=np.where(expZpNCCD > 0)
        expZpMean[gd] /= expZpNCCD[gd]

        fig=plt.figure(1,figsize=(8,6))
        fig.clf()

        ax=fig.add_subplot(111)

        firstMJD = np.floor(np.min(self.fgcmPars.expMJD))

        # FIXME: make configurable
        cols = ['g','r','b','m','y']
        syms = ['.','+','o','*','x']

        for i in xrange(self.fgcmPars.nBands):
            use,=np.where((self.fgcmPars.expBandIndex == i) &
                          (expZpMean > 0.0))

            if (use.size == 0) :
                continue

            plt.plot(self.fgcmPars.expMJD[use] - firstMJD,
                     expZpMean[use],cols[i]+syms[i],
                     label=r'$(%s)$' % (self.fgcmPars.bands[i]))

        ax.legend(loc=3)

        fig.savefig('%s/%s_zeropoints.png' % (self.plotPath,
                                              self.outfileBaseWithCycle))
        plt.close(fig)

    def _computeZpt(self,zpStruct,zpIndex):
        """
        Internal method to compute the zeropoint from constituents

        parameters
        ----------
        zpStruct: recarray
           Zero point structure
        zpIndex: int array
           Array of indices to compute zeropoints
        """

        zpStruct['FGCM_ZPT'][zpIndex] = (2.5*np.log10(zpStruct['FGCM_I0'][zpIndex]) +
                                         zpStruct['FGCM_FLAT'][zpIndex] +
                                         zpStruct['FGCM_DUST'][zpIndex] +
                                         zpStruct['FGCM_APERCORR'][zpIndex] +
                                         2.5*np.log10(zpStruct['EXPTIME'][zpIndex]) +
                                         self.zptAB +
                                         zpStruct['FGCM_GRY'][zpIndex])

    def _computeZptErr(self,zpStruct,zpExpIndex,zpIndex):
        """
        Internal method to compute zeropoint error from constituents

        parameters
        ----------
        zpStruct: recarray
           Zero point structure
        zpExpIndex: int array
           Index to go from zeropoint structure to exposures
        zpIndex: int array
           Array of indices to compute zeropoint errors
        """

        # sigFgcm is computed per *band* not per *filter*
        sigFgcm = self.fgcmPars.compSigFgcm[self.fgcmPars.expBandIndex[zpExpIndex[zpIndex]]]
        nTilingsM1 = np.clip(zpStruct['FGCM_TILINGS'][zpIndex]-1.0,1.0,1e10)

        zpStruct['FGCM_ZPTERR'][zpIndex] = np.sqrt((sigFgcm**2./nTilingsM1) +
                                                   zpStruct['FGCM_ZPTVAR'][zpIndex] +
                                                   self.sigma0Cal**2.)

    def saveZptFits(self):
        """
        Save zeropoint to fits file
        """

        import fitsio

        outFile = '%s/%s_zpt.fits' % (self.outputPath,self.outfileBaseWithCycle)
        self.fgcmLog.info('Saving zeropoints to %s' % (outFile))
        fitsio.write(outFile,self.zpStruct,clobber=True,extname='ZPTS')

    def saveAtmFits(self):
        """
        Save atmosphere parameters to fits file
        """

        import fitsio

        outFile = '%s/%s_atm.fits' % (self.outputPath,self.outfileBaseWithCycle)
        self.fgcmLog.info('Saving atmosphere parameters to %s' % (outFile))
        fitsio.write(outFile,self.atmStruct,clobber=True,extname='ATMPARS')

class FgcmZeropointPlotter(object):
    """
    Class to make zeropoint plots

    parameters
    ----------
    zpStruct: recarray
       Zero point structure
    filterNames: string array
       Names of filters
    plotPath: string
       Directory to make plots
    outfileBase: string
       Output file base string
    """

    def __init__(self, zpStruct, filterNames, plotPath, outfileBase):
        self.zpStruct = zpStruct
        self.filterNames = filterNames
        self.plotPath = plotPath
        self.outfileBase = outfileBase

    def makeR1I1Plots(self):
        """
        Make R1 vs I1 plots.
        """

        acceptMask = (zpFlagDict['PHOTOMETRIC_FIT_EXPOSURE'] |
                      zpFlagDict['PHOTOMETRIC_EXTRA_EXPOSURE'])
        for filterName in self.filterNames:
            use,=np.where((np.core.defchararray.rstrip(self.zpStruct['FILTERNAME']) == filterName.encode('utf-8')) &
                          ((self.zpStruct['FGCM_FLAG'] & acceptMask) > 0) &
                          (np.abs(self.zpStruct['FGCM_R10']) < 1000.0) &
                          (np.abs(self.zpStruct['FGCM_R0']) < 1000.0))

            if (use.size == 0):
                continue

            i1 = self.zpStruct['FGCM_I10'][use] * self.zpStruct['FGCM_I0'][use]
            r1 = self.zpStruct['FGCM_R10'][use] * self.zpStruct['FGCM_R0'][use]

            # limit to a reasonable range
            #  note that r1 is much noisier than i1
            ok, = np.where((r1 > (i1.min()-2.0)) &
                           (r1 < (i1.max()+2.0)))

            i1=i1[ok]
            r1=r1[ok]

            fig = plt.figure(1,figsize=(8,6))
            fig.clf()

            ax=fig.add_subplot(111)

            ax.hexbin(i1,r1,cmap=plt.get_cmap('gray_r'),rasterized=True)
            # and overplot a 1-1 line that best covers the range of the data
            xlim = ax.get_xlim()
            range0 = xlim[0]+0.001
            range1 = xlim[1]-0.001
            ax.plot([range0,range1],[range0,range1],'b--',linewidth=2)

            ax.set_xlabel(r'$I_1$ from FGCM Fit',fontsize=16)
            ax.set_ylabel(r'$R_1$ from Retrieval',fontsize=16)

            text=r'$(%s)$' % (filterName)
            ax.annotate(text,(0.1,0.93),xycoords='axes fraction',
                        ha='left',va='top',fontsize=16)

            fig.savefig('%s/%s_i1r1_%s.png' % (self.plotPath,
                                               self.outfileBase,
                                               filterName))
            plt.close(fig)

    def makeR1I1Maps(self, ccdOffsets, ccdField='CCDNUM'):
        """
        Make R1, I1, and R1-I1 maps.

        parameters
        ----------
        ccdOffsets: ccd offset struct
        ccdField: string, default='CCDNUM'
        """

        from .fgcmUtilities import plotCCDMap

        acceptMask = (zpFlagDict['PHOTOMETRIC_FIT_EXPOSURE'] |
                      zpFlagDict['PHOTOMETRIC_EXTRA_EXPOSURE'])

        plotTypes=['I1', 'R1', 'R1 - I1']

        ccdMin = np.min(self.zpStruct[ccdField])
        ccdMax = np.max(self.zpStruct[ccdField])
        nCCD = (ccdMax - ccdMin) + 1

        for filterName in self.filterNames:
            use0,=np.where((np.core.defchararray.rstrip(self.zpStruct['FILTERNAME']) == filterName.encode('utf-8')) &
                           ((self.zpStruct['FGCM_FLAG'] & acceptMask) > 0) &
                           (np.abs(self.zpStruct['FGCM_R10']) < 1000.0) &
                           (np.abs(self.zpStruct['FGCM_R0']) < 1000.0))

            if (use0.size == 0):
                continue

            ccdIndex = np.searchsorted(np.arange(ccdMin,ccdMax+1),
                                       self.zpStruct[ccdField][use0])

            i1 = self.zpStruct['FGCM_I10'][use0] * self.zpStruct['FGCM_I0'][use0]
            r1 = self.zpStruct['FGCM_R10'][use0] * self.zpStruct['FGCM_R0'][use0]

            meanI1 = np.zeros(nCCD)
            meanR1 = np.zeros(nCCD)
            nPerCCD = np.zeros(nCCD,dtype=np.int32)

            np.add.at(meanI1, ccdIndex, i1)
            np.add.at(meanR1, ccdIndex, r1)
            np.add.at(nPerCCD, ccdIndex, 1)

            use,=np.where(nPerCCD > 0)
            meanI1[use] /= nPerCCD[use]
            meanR1[use] /= nPerCCD[use]

            # use the same range scale for all the plots
            st = np.argsort(meanR1[use])
            lo = meanR1[use[st[int(0.02*st.size)]]]
            hi = meanR1[use[st[int(0.98*st.size)]]]

            for plotType in plotTypes:
                fig=plt.figure(1,figsize=(8,6))
                fig.clf()

                ax=fig.add_subplot(111)

                if (plotType == 'R1'):
                    plotCCDMap(ax, ccdOffsets[use], meanR1[use], plotType, loHi=[lo,hi])
                elif (plotType == 'I1'):
                    plotCCDMap(ax, ccdOffsets[use], meanI1[use], plotType, loHi=[lo,hi])
                else:
                    # for the residuals, center at zero, but use lo/hi
                    amp = np.abs((hi - lo)/2.)
                    plotCCDMap(ax, ccdOffsets[use], meanR1[use] - meanI1[use], plotType, loHi=[-amp, amp])

                text = r'$(%s)$' % (filterName) + '\n' + \
                    r'%s' % (plotType)
                ax.annotate(text,
                            (0.1,0.93),xycoords='axes fraction',
                            ha='left',va='top',fontsize=18)

                fig.savefig('%s/%s_%s_%s.png' % (self.plotPath,
                                                 self.outfileBase,
                                                 plotType.replace(" ",""),
                                                 filterName))
                plt.close(fig)

        return None


# -*- coding: utf-8 -*-
from enum import Enum
from typing import List, Tuple, Union
from numpy import *
import numpy as np
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
from k_means import k_means, get_group_center
from ransac import *
from blend import *
from L2_Net import L2Net
import time

def show_image(image: np.ndarray) -> None:
    from PIL import Image
    Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).show()

class Method(Enum):

    SURF = cv2.xfeatures2d.SURF_create
    SIFT = cv2.xfeatures2d.SIFT_create
    ORB = cv2.ORB_create

colors = ((123, 234, 12), (23, 44, 240), (224, 120, 34), (21, 234, 190),
          (80, 160, 200), (243, 12, 100), (25, 90, 12), (123, 10, 140))

class Area:

    def __init__(self, *points):

        self.points = list(points)

    def is_inside(self, x: Union[float, Tuple[float, float]], y: float=None):
        if isinstance(x, tuple):
            x, y = x
        raise NotImplementedError()

class Matcher:
    def __init__(self, image1: np.ndarray, image2: np.ndarray, method: Enum = Method.SIFT, threshold=800, ratio=400):
        """输入两幅图像，计算其特征值
        此类用于输入两幅图像，计算其特征值，输入两幅图像分别为numpy数组格式的图像，
        其中的method参数要求输入SURF、SIFT或者ORB，threshold参数为特征值检测所需的阈值。

        Args:
            image1 (np.ndarray): 图像一
            image2 (np.ndarray): 图像二
            method (Enum, optional): Defaults to Method.SIFT. 特征值检测方法
            ratio (int, optional): Defaults to 400. L2-Net特征向量比重
            threshold (int, optional): Defaults to 800. 特征值阈值

        """

        self.image1 = image1
        self.image2 = image2
        self.method = method
        self.ratio = ratio

        self.ros1 = []
        # type np.ndarray
        self.ros1_3 = None
        self.ros2 = []
        # type np.ndarray
        self.ros2_3 = None
        self.loc1 = []
        # type list
        self.loc2 = []
        # type list
        self.threshold = threshold

        self._keypoints1 = None
        # type List[cv2.KeyPoint]
        self._descriptors1 = None
        # type np.ndarray
        self._keypoints2 = None
        # type List[cv2.KeyPoint]
        self._descriptors2 = None
        # type np.ndarray

        if self.method == Method.ORB:
            # error if not set this
            self.matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
        else:
            self.matcher = cv2.BFMatcher(cv2.NORM_L2, crossCheck=False)
            # self.matcher = cv2.FlannBasedMatcher()

        self.match_points = []
        self.match_points1 = []

        self.image_points1 = np.array([])
        self.image_points2 = np.array([])

        self.keypoints1_loc = None
        self.keypoints2_loc = None

        self.x = 0
        self.y = 0

        self.vggd1 = []
        self.vggd2 = []

        self.l2_1des = None
        self.l2_2des = None

        self.save1 = []
        self.save2 = []
        self.save3 = []

    def compute_keypoint(self) -> None:
        """计算特征点
        利用给出的特征值检测方法对图像进行特征值检测。

        Args:
            image (np.ndarray): 图像
        """
        print('Compute keypoint by SIFT')
        print(Method.SIFT)

        feature = self.method.value(self.threshold)
        self._keypoints1, self._descriptors1 = feature.detectAndCompute(
            self.image1, None)
        self._keypoints2, self._descriptors2 = feature.detectAndCompute(
            self.image2, None)
        self.keypoints1_loc = cv2.KeyPoint_convert(self._keypoints1)
        self.keypoints2_loc = cv2.KeyPoint_convert(self._keypoints2)

    def get_the_ros(self):
        """选取待匹配的特征块
                将灰度图像选取特定区域，并转换维度为(?, 32, 32, 1)
                """

        for i in range(len(self.keypoints1_loc)):
            self.loc1.append(cv2.KeyPoint(x=self.keypoints1_loc[i][0], y=self.keypoints1_loc[i][1], _size=32,
                                          _angle=-1, _response=0.018, _octave=1, _class_id=-1))
            if self.keypoints1_loc[i][0] < 16:
                self.x = 16
            elif self.keypoints1_loc[i][0] > 240:
                self.x = 240
            else:
                self.x = int(self.keypoints1_loc[i][0])

            if self.keypoints1_loc[i][1] < 16:
                self.y = 16
            elif self.keypoints1_loc[i][1] > 240:
                self.y = 240
            else:
                self.y = int(self.keypoints1_loc[i][1])
            self.ros1.append(self.image1[self.x-16: self.x+16, self.y-16: self.y+16])

        for i in range(len(self.keypoints2_loc)):
            self.loc2.append(cv2.KeyPoint(x=self.keypoints2_loc[i][0], y=self.keypoints2_loc[i][1], _size=32,
                                          _angle=-1, _response=0.018, _octave=1, _class_id=-1))
            if self.keypoints2_loc[i][0] < 16:
                self.x = 16
            elif self.keypoints2_loc[i][0] > 240:
                self.x = 240
            else:
                self.x = int(self.keypoints2_loc[i][0])

            if self.keypoints2_loc[i][1] < 16:
                self.y = 16
            elif self.keypoints2_loc[i][1] > 240:
                self.y = 240
            else:
                self.y = int(self.keypoints2_loc[i][1])
            self.ros2.append(self.image2[self.x-16: self.x+16, self.y-16: self.y+16])

        # self.ros1, self.loc1 = slice(self.image1)
        # self.ros2, self.loc2 = slice(self.image2)

    def compute_kepoint_by_L2_Net(self) -> None:
        """
        通过VGG计算描述向量+PCA
        筛选并存储描述向量和keypoint。

        """
        self.get_the_ros()
        self.ros1 = np.array(self.ros1)
        self.ros2 = np.array(self.ros2)

        print('Compute keypoint by L2-net')
        l2net = L2Net("L2Net-HP+")

        self.l2_1des = l2net.calc_descriptors(self.ros1)
        self.l2_2des = l2net.calc_descriptors(self.ros2)

    def match(self, max_match_lenth=200, threshold=0.04, show_match=False):
        """对两幅图片计算得出的特征值进行匹配，对ORB来说使用OpenCV的BFMatcher算法，而对于其他特征检测方法则使用FlannBasedMatcher算法。

            max_match_lenth (int, optional): Defaults to 20. 最大匹配点数量
            threshold (float, optional): Defaults to 0.04. 默认最大匹配距离差
            show_match (bool, optional): Defaults to False. 是否展示匹配结果
        """

        self.compute_keypoint()
        self.compute_kepoint_by_L2_Net()
        good = []

        '''计算两张图片中的配对点，并至多取其中最优的`max_match_lenth`个'''

        self.match_points = self.matcher.knnMatch(self._descriptors1, self._descriptors2, k=2)

        for m, n in self.match_points:

            if m.distance < 0.6 * n.distance:
                good.append(m)

        print(len(good))

        ptsA = np.float32([self._keypoints1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
        ptsB = np.float32([self._keypoints2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
        ransacReprojThreshold = 4
        H, status = cv2.findHomography(ptsA, ptsB, cv2.RANSAC, ransacReprojThreshold)

        matchesMask = status.ravel().tolist()

        h, w = 256, 256
        pts = np.float32([[0, 0], [0, h - 1], [w - 1, h - 1], [w - 1, 0]]).reshape(-1, 1, 2)
        dst = cv2.perspectiveTransform(pts, H)

        img2 = cv2.polylines(self.image2, [np.int32(dst)], True, 255, 3, cv2.LINE_AA)
        draw_params = dict(matchColor=(0, 255, 0),
                           singlePointColor=None,
                           matchesMask=matchesMask,
                           flags=2)
        img3 = cv2.drawMatches(self.image1, self._keypoints1, img2, self._keypoints2, good, None, **draw_params)
        show_image(img3)
        '''由最佳匹配取得匹配点对，并进行形变拼接'''

def get_weighted_points(image_points: np.ndarray):

    average = np.average(image_points, axis=0)

    max_index = np.argmax(np.linalg.norm((image_points - average), axis=1))
    return np.append(image_points, np.array([image_points[max_index]]), axis=0)

def pca(image, rate=128):
    mean_value = mean(image, axis=0)
    image = image-mean_value
    c = cov(image, rowvar=False)
    eigvalue, eigvector = linalg.eig(mat(c))

    index_vec = np.argsort(-eigvalue)
    n_largest_index = index_vec[:rate]
    t = eigvector[:, n_largest_index]
    print(np.shape(image), np.shape(t))
    new_image = np.dot(image, t)
    return new_image, t, mean_value

class Stitcher:

    def __init__(self, image1: np.ndarray, image2: np.ndarray, method: Enum = Method.SIFT, use_kmeans=False):
        """输入图像和匹配，对图像进行拼接
        目前采用简单矩阵匹配和平均值拼合

        Args:
            image1 (np.ndarray): 图像一
            image2 (np.ndarray): 图像二
            matcher (Matcher): 匹配结果
            use_kmeans (bool): 是否使用kmeans 优化点选择
        """
        self.image_points1, self.image_points2 = None, None
        self.image1 = image1
        self.image2 = image2
        self.method = method
        self.use_kmeans = use_kmeans
        self.matcher = Matcher(image1, image2, method=method)
        self.M = np.eye(3)

        self.image = None

    def stich(self, show_result=True, max_match_lenth=40, show_match_point=True, use_partial=False,
              use_new_match_method=False, use_gauss_blend=True):
        """对图片进行拼合

            show_result (bool, optional): Defaults to True. 是否展示拼合图像
            show_match_point (bool, optional): Defaults to True. 是否展示拼合点
        """
        self.matcher.match(max_match_lenth=max_match_lenth,
                           show_match=show_match_point)

        if self.use_kmeans:
            self.image_points1, self.image_points2 = get_group_center(
                self.matcher.image_points1, self.matcher.image_points2)
        else:
            self.image_points1, self.image_points2 = (
                self.matcher.image_points1, self.matcher.image_points2)

        if use_new_match_method:
            self.M = GeneticTransform(self.image_points1, self.image_points2).run()
        else:
            self.M, _ = cv2.findHomography(self.image_points1, self.image_points2, method=cv2.RANSAC)

        print("Good points and average distance: ", GeneticTransform.get_value(
            self.image_points1, self.image_points2, self.M))

        left, right, top, bottom = self.get_transformed_size()
        width = int(max(right, self.image2.shape[1]) - min(left, 0))
        height = int(max(bottom, self.image2.shape[0]) - min(top, 0))

        if use_partial:
            self.partial_transform()

        # 移动矩阵
        self.adjustM = np.array(
            [[1, 0, max(-left, 0)],  # 横向
             [0, 1, max(-top, 0)],  # 纵向
             [0, 0, 1]
             ], dtype=np.float64)
        self.M = np.dot(self.adjustM, self.M)
        transformed_1 = cv2.warpPerspective(
            self.image1, self.M, (width, height))
        transformed_2 = cv2.warpPerspective(
            self.image2, self.adjustM, (width, height))

        self.image = self.blend(transformed_1, transformed_2, use_gauss_blend=use_gauss_blend)

        if show_match_point:
            for point1, point2 in zip(self.image_points1, self.image_points2):
                point1 = self.get_transformed_position(tuple(point1))
                point1 = tuple(map(int, point1))
                point2 = self.get_transformed_position(tuple(point2), M=self.adjustM)
                point2 = tuple(map(int, point2))
                cv2.circle(self.image, point1, 10, (20, 20, 255), 5)
                cv2.circle(self.image, point2, 8, (20, 200, 20), 5)

    def blend(self, image1: np.ndarray, image2: np.ndarray, use_gauss_blend=True) -> np.ndarray:
        """对图像进行融合

        Args:
            image1 (np.ndarray): 图像一
            image2 (np.ndarray): 图像二
            use_gauss_blend
        Returns:
            np.ndarray: 融合结果
        """

        mask = self.generate_mask(image1, image2)
        print("Blending")
        if use_gauss_blend:
            result = gaussian_blend(image1, image2, mask, mask_blend=10)
        else:
            result = direct_blend(image1, image2, mask, mask_blend=0)

        return result

    def generate_mask(self, image1: np.ndarray, image2: np.ndarray):
        """生成供融合使用的遮罩，由变换后图像的垂直平分线来构成分界线

        Args:
            image1
            image2

        Returns:
            np.ndarray: 01数组
        """
        print("Generating mask")
        # x, y
        center1 = self.image1.shape[1] / 2, self.image1.shape[0] / 2
        center1 = self.get_transformed_position(center1)
        center2 = self.image2.shape[1] / 2, self.image2.shape[0] / 2
        center2 = self.get_transformed_position(center2, M=self.adjustM)
        x1, y1 = center1
        x2, y2 = center2

        def function(y, x, *z):
            return (y2 - y1) * y < -(x2 - x1) * (x - (x1 + x2) / 2) + (y2 - y1) * (y1 + y2) / 2

        mask = np.fromfunction(function, image1.shape)

        mask = np.logical_and(mask, np.logical_not(image2)) \
            + np.logical_and(mask, image1)\
            + np.logical_and(image1, np.logical_not(image2))

        return mask

    def get_transformed_size(self) -> Tuple[int, int, int, int]:
        """计算形变后的边界
        计算形变后的边界，从而对图片进行相应的位移，保证全部图像都出现在屏幕上。

        Returns:
            Tuple[int, int, int, int]: 分别为左右上下边界
        """

        conner_0 = (0, 0)
        conner_1 = (self.image1.shape[1], 0)
        conner_2 = (self.image1.shape[1], self.image1.shape[0])
        conner_3 = (0, self.image1.shape[0])
        points = [conner_0, conner_1, conner_2, conner_3]

        top = min(map(lambda x: self.get_transformed_position(x)[1], points))
        bottom = max(
            map(lambda x: self.get_transformed_position(x)[1], points))
        left = min(map(lambda x: self.get_transformed_position(x)[0], points))
        right = max(map(lambda x: self.get_transformed_position(x)[0], points))

        return left, right, top, bottom

    def get_transformed_position(self, x: Union[float, Tuple[float, float]], y: float = None, M = None) -> Tuple[float, float]:
        """求得某点在变换矩阵（self.M）下的新坐标

        Args:
            x (Union[float, Tuple[float, float]]): x坐标或(x,y)坐标
            y (float, optional): Defaults to None. y坐标，可无
            M (np.ndarray, optional): Defaults to None. 利用M进行坐标变换运算

        Returns:
            Tuple[float, float]:  新坐标
        """

        if isinstance(x, tuple):
            x, y = x
        p = np.array([x, y, 1])[np.newaxis].T
        if M is not None:
            M = M
        else:
            M = self.M
        pa = np.dot(M, p)
        return pa[0, 0] / pa[2, 0], pa[1, 0] / pa[2, 0]

if __name__ == "__main__":

    os.chdir(os.path.dirname(__file__))

    start_time = time.time()
    img1 = cv2.imread("./example/label.jpg", cv2.IMREAD_GRAYSCALE)
    img2 = cv2.imread("./example/4.jpg", cv2.IMREAD_GRAYSCALE)
    img1 = np.expand_dims(img1, axis=2)
    img2 = np.expand_dims(img2, axis=2)

    stitcher = Stitcher(img1, img2, Method.SIFT, False)
    stitcher.stich(max_match_lenth=50, use_partial=False, use_new_match_method=True, use_gauss_blend=False)

    cv2.imwrite('./example/merge.jpg', stitcher.image)

    print("Time: ", time.time() - start_time)
    print("M: ", stitcher.M)


# from __future__ import absolute_import
# from __future__ import division
# from __future__ import print_function

import argparse
import os
import pprint

import torch
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.nn.functional as f
# import torch.optim
# import torch.utils.data
# import torch.utils.data.distributed
import torchvision.transforms as transforms
from pose_hrnet import get_pose_net
# import coremltools as ct
from collections import OrderedDict
from config import cfg
from config import update_config

from PIL import Image
import numpy as np
import cv2

from utils import pose_process, plot_pose
from natsort import natsorted

mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)

index_mirror = np.concatenate([
                [1,3,2,5,4,7,6,9,8,11,10,13,12,15,14,17,16],
                [21,22,23,18,19,20],
                np.arange(40,23,-1), np.arange(50,40,-1),
                np.arange(51,55), np.arange(59,54,-1),
                [69,68,67,66,71,70], [63,62,61,60,65,64],
                np.arange(78,71,-1), np.arange(83,78,-1),
                [88,87,86,85,84,91,90,89],
                np.arange(113,134), np.arange(92,113)
                ]) - 1
assert(index_mirror.shape[0] == 133)

multi_scales = [512,640]
def norm_numpy_totensor(img):
    img = img.astype(np.float32) / 255.0
    for i in range(3):
        img[:, :, :, i] = (img[:, :, :, i] - mean[i]) / std[i]
    return torch.from_numpy(img).permute(0, 3, 1, 2)
def stack_flip(img):
    img_flip = cv2.flip(img, 1)
    return np.stack([img, img_flip], axis=0)

def merge_hm(hms_list):
    assert isinstance(hms_list, list)
    for hms in hms_list:
        hms[1,:,:,:] = torch.flip(hms[1,index_mirror,:,:], [2])

    hm = torch.cat(hms_list, dim=0)
    # print(hm.size(0))
    hm = torch.mean(hms, dim=0)
    return hm
def main():

    with torch.no_grad():
        config = 'wholebody_w48_384x288.yaml'
        cfg.merge_from_file(config)

        # dump_input = torch.randn(1, 3, 256, 256)
        # newmodel = PoseHighResolutionNet()
        newmodel = get_pose_net(cfg, is_train=False)
        print(newmodel)
        # dump_output = newmodel(dump_input)
        # print(dump_output.size())
        checkpoint = torch.load('./hrnet_w48_coco_wholebody_384x288-6e061c6a_20200922.pth')
        # newmodel.load_state_dict(checkpoint['state_dict'])

        state_dict = checkpoint['state_dict']
        new_state_dict = OrderedDict()
        for k, v in state_dict.items():
            if 'backbone.' in k:
                name = k[9:] # remove module.
            if 'keypoint_head.' in k:
                name = k[14:] # remove module.
            new_state_dict[name] = v
        newmodel.load_state_dict(new_state_dict)

        newmodel.cuda().eval()

        transform  = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
        ])

        input_path = '/home/ubuntu/workspace/data/sign/val'
        paths = []
        names = []
        for root, _, fnames in natsorted(os.walk(input_path)):
            for fname in natsorted(fnames):
                path1 = os.path.join(root, fname)
                if 'depth' in fname:
                    continue
                paths.append(path1)
                names.append(fname)
        print(len(paths))
        # paths = paths[:4]
        # names = names[:4]
        step = 600
        start_step = 6
        # paths = paths[start_step*step:(start_step+1)*step]
        # names = names[start_step*step:(start_step+1)*step]
        paths = paths[4200:]
        names = names[4200:]
        paths = paths[::-1]
        names = names[::-1]

        for i, path in enumerate(paths):
            # if i > 1:
            #     break
            output_npy = 'npy3/{}.npy'.format(names[i])

            if os.path.exists(output_npy):
                continue

            cap = cv2.VideoCapture(path)

            frame_width = int(cap.get(3))
            frame_height = int(cap.get(4))
            # frame_width = 256
            # frame_height = 256
            print(path)
            # output_filename = os.path.join('out_test', names[i])

            # img = Image.open(image_path)
            # fps = cap.get(cv2.CAP_PROP_FPS)
            # writer = cv2.VideoWriter(output_filename,cv2.VideoWriter_fourcc('M','P','4','V'), 5, (frame_width,frame_height))
            output_list = []

            while cap.isOpened():
                success, img = cap.read()
                if not success:
                    print("Ignoring empty camera frame.")
                    # If loading a video, use 'break' instead of 'continue'.
                    break
                # img = cv2.resize(img, (256,256))
                frame_height, frame_width = img.shape[:2]
                img = cv2.flip(img, flipCode=1)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                out = []
                for scale in multi_scales:
                    if scale != 512:
                        img_temp = cv2.resize(img, (scale,scale))
                    else:
                        img_temp = img
                    img_temp = stack_flip(img_temp)
                    img_temp = norm_numpy_totensor(img_temp).cuda()
                    hms = newmodel(img_temp)
                    if scale != 512:
                        out.append(f.interpolate(hms, (frame_width // 4,frame_height // 4), mode='bilinear'))
                    else:
                        out.append(hms)

                out = merge_hm(out)
                # print(out.size())
                # hm, _ = torch.max(out, 1)
                # hm = hm.cpu().numpy()
                # print(hm.shape)
                # np.save('hm.npy', hm)
                result = out.reshape((133,-1))
                result = torch.argmax(result, dim=1)
                # print(result)
                result = result.cpu().numpy().squeeze()

                # print(result.shape)
                y = result // (frame_width // 4)
                x = result % (frame_width // 4)
                pred = np.zeros((133, 3), dtype=np.float32)
                pred[:, 0] = x
                pred[:, 1] = y

                hm = out.cpu().numpy().reshape((133, frame_height//4, frame_height//4))

                pred = pose_process(pred, hm)
                pred[:,:2] *= 4.0
                # print(pred.shape)
                assert pred.shape == (133, 3)

                # print(arg.cpu().numpy())
                # np.save('npy/{}.npy'.format(names[i]), np.array([x,y,score]).transpose())
                output_list.append(pred)
                # img = np.asarray(img)
                # for j in range(133):
                #     img = cv2.circle(img, (int(x[j]), int(y[j])), radius=2, color=(255,0,0), thickness=-1)
                # img = plot_pose(img, pred)
                # cv2.imwrite('out/{}.png'.format(names[i]), cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
                # writer.write(cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
            output_list = np.array(output_list)
            # print(output_list.shape)
            np.save(output_npy, output_list)
            cap.release()
            # writer.release()
            # break

if __name__ == '__main__':
    main()


import os
import json
import argparse
import torch
import random
import torch.optim as optim
from model_vc import StyleEncoder
from model_vc import Generator
from dataset import AudiobookDataset
from dataset import train_collate
from dataset import test_collate
from utils.dsp import save_wav
import numpy as np
from hparams import hparams as hp

def save_checkpoint(device, model, optimizer, checkpoint_dir, epoch):
    checkpoint_path = os.path.join(
        checkpoint_dir, "checkpoint_step{:06d}.pth".format(epoch))
    optimizer_state = optimizer.state_dict()
    torch.save({
        "model": model.state_dict(),
        "optimizer": optimizer_state,
        "epoch": epoch
    }, checkpoint_path)
    print("Saved checkpoint:", checkpoint_path)

def load_checkpoint(path, model, device, optimizer, reset_optimizer=False):
    print("Load checkpoint from: {}".format(path))
    checkpoint = torch.load(path, map_location=device)
    model.load_state_dict(checkpoint["model"])
    if not reset_optimizer:
        optimizer_state = checkpoint["optimizer"]
        if optimizer_state is not None:
            print("Load optimizer state from {}".format(path))
            optimizer.load_state_dict(checkpoint["optimizer"])
    epoch = checkpoint['epoch']
    return epoch

def train(args, model, device, train_loader, optimizer, epoch, sigma=1.0):
    model.train()
    train_loss = 0

    for batch_idx, (m, e) in enumerate(train_loader):
        m = m.to(device)
        e = e.to(device)

        model.zero_grad()

        mel_outputs, mel_outputs_postnet, codes = model(m, e, e)

        m_rec = mel_outputs_postnet
        codes_rec = model(m_rec, e, None)

        L_recon = ((mel_outputs_postnet - m) ** 2).sum(dim=(1,2)).mean()
        L_recon0 = ((mel_outputs - m) ** 2).sum(dim=(1,2)).mean()
        L_content = torch.abs(codes - codes_rec).sum(dim=1).mean()

        loss = L_recon + L_recon0 + L_content

        loss.backward()
        optimizer.step()

        train_loss += loss.item() * len(m)

        if batch_idx % 10 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(m), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

    train_loss /= len(train_loader.dataset)
    print('\nTrain set: Average loss: {:.4f}\n'.format(train_loss))

def test(model, device, test_loader, checkpoint_dir, epoch, sigma=1.0):
    print("Using averaged model for evaluation")
    model.eval()

    test_loss = 0

    with torch.no_grad():
        for batch_idx, (m, e) in enumerate(test_loader):
            m = m.to(device)
            e = e.to(device)

            mel_outputs, mel_outputs_postnet, codes = model(m, e, e)

            m_rec = mel_outputs_postnet
            codes_rec = model(m_rec, e, None)

            L_recon = ((mel_outputs_postnet - m) ** 2).sum(dim=(1,2)).mean()
            L_recon0 = ((mel_outputs - m) ** 2).sum(dim=(1,2)).mean()
            L_content = torch.abs(codes - codes_rec).sum(dim=1).mean()

            loss = L_recon + L_recon0 + L_content

            if batch_idx % 100 == 0:
                print('Val Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, batch_idx * len(m), len(test_loader.dataset),
                    100. * batch_idx / len(test_loader), loss.item()))
            test_loss += loss.item()

        test_loss /= len(test_loader.dataset)
        print('\nTest set: Average loss: {:.4f}\n'.format(test_loss))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Train or run some neural net')
    parser.add_argument('-d', '--data', type=str, default='./data', help='dataset directory')
    parser.add_argument('--checkpoint', type=str, default=None,
                        help='The path to checkpoint')
    parser.add_argument('--epochs', type=int, default=600,
                        help='number of epochs to train (default: 14)')
    parser.add_argument('--batch-size', type=int, default=8, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--lr', type=float, default=1e-4, metavar='LR',
                        help='learning rate (default: 1.0)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    args = parser.parse_args()

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

    data_path = args.data

    use_cuda = not args.no_cuda and torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")

    kwargs = {'num_workers': 8, 'pin_memory': True} if use_cuda else {}

    torch.autograd.set_detect_anomaly(True)

    with open(os.path.join(data_path, 'train_data.json'), 'r') as f:
        train_data = json.load(f)

    with open(os.path.join(data_path, 'test_data.json'), 'r') as f:
        test_data = json.load(f)

    train_loader = torch.utils.data.DataLoader(
        AudiobookDataset(train_data),
        collate_fn=train_collate,
        batch_size=args.batch_size, shuffle=True, **kwargs)

    test_loader = torch.utils.data.DataLoader(
        AudiobookDataset(test_data),
        collate_fn=test_collate,
        batch_size=1, shuffle=False, **kwargs)

    model = Generator(hp.dim_neck, hp.dim_emb, hp.dim_pre, hp.freq).to(device)
    optimizer = optim.Adam(model.parameters(), lr=args.lr)

    current_epoch = 0
    if args.checkpoint:
        current_epoch = load_checkpoint(args.checkpoint, model, device, optimizer)

    checkpoint_dir = 'checkpoints'
    os.makedirs(checkpoint_dir, exist_ok=True)

    for epoch in range(current_epoch + 1, args.epochs + 1):
        print(f'epoch {epoch}')
        train(args, model, device, train_loader, optimizer, epoch)

        if epoch % 10 == 0:
            test(model, device, test_loader, checkpoint_dir, epoch)
            save_checkpoint(device, model, optimizer, checkpoint_dir, epoch)


# coding=utf8

import cv2
import numpy as np
from cut import cut
from nms import nms
from rerec import rerec
from bbreg import bbreg
from loader import PNetPredictor, RNetPredictor, ONetPredictor

class MTCNN(object):

    def __init__(self, pnet_model_path, rnet_model_path, onet_model_path):

        self.pnet = PNetPredictor(pnet_model_path)
        self.rnet = RNetPredictor(rnet_model_path)
        self.onet = ONetPredictor(onet_model_path)

    def detect(self, img, min_size=40, factor=0.709,
               score_threshold=[0.5, 0.5, 0.5],
               iou_threshold=[0.5, 0.7, 0.7, 0.7]):
        '''

        :param img: numpy array (shape [height, width, 3] )
        :return: boxes: numpy array (shape [N, 5])
                                [h0, w0, h1, w1, score]
                 landmarks: numpy array (shape [N, 10]
                            [h_leye, h_reye, h_nose, h_lmouth, h_rmouth,
                             w_leye, w_reye, w_nose, w_lmouth, w_rmouth]
        '''
        # Stage 1
        h, w, _ = img.shape
        min_l = min(h, w)
        m = 12. / min_size
        min_l = min_l * m
        scales = []
        while min_l >= 12:
            scales.append(m)
            min_l *= factor
            m *= factor

        proposals_total = np.zeros((0, 9), dtype=np.float32)
        for scale in scales:
            hs = int(np.round(h * scale))
            ws = int(np.round(w * scale))
            img_resize = cv2.resize(
                img, (ws, hs), interpolation=cv2.INTER_LINEAR
            )
            img_resize = (img_resize - 127.5) * 0.0078125
            img_resize = img_resize[np.newaxis, :, :, :]
            try:
                cls_prob, bbox_reg = self.pnet.predict(img_resize)
            except:
                continue
            prob = cls_prob[..., 1]
            idx_pos = np.where(prob >= score_threshold[0])
            prob_pos = prob[idx_pos].reshape([-1, 1])
            h0 = idx_pos[1].reshape([-1, 1])
            w0 = idx_pos[2].reshape([-1, 1])
            bbox_pos = np.concatenate((h0, w0, h0+5.5, w0+5.5), axis=1)
            bbox_pos = np.round(bbox_pos * 2 / scale)
            bbox_reg_pos = bbox_reg[idx_pos]
            proposals = np.concatenate(
                (bbox_pos, prob_pos, bbox_reg_pos), axis=1
            )
            proposals = proposals[nms(proposals, iou_threshold[0])]
            proposals_total = np.concatenate(
                (proposals_total, proposals), axis=0
            )
        if len(proposals_total) == 0:
            return (np.zeros((0, 5), dtype=np.float32),
                    np.zeros((0, 10), dtype=np.float32))

        proposals_total = proposals_total[nms(proposals_total,
                                          threshold=iou_threshold[1])]
        proposals_total = bbreg(proposals_total)

        # Stage 2
        proposals_total = rerec(proposals_total)
        boxes = proposals_total[:, 0:4].astype(np.int32)
        input = cut(img, boxes, h, w, 24)
        input = (input - 127.5) * 0.0078125
        try:
            cls_prob, bbox_reg = self.rnet.predict(input)
        except:
            return (np.zeros((0, 5), dtype=np.float32),
                    np.zeros((0, 10), dtype=np.float32))

        prob = cls_prob[..., 1]
        idx_pos = np.where(prob >= score_threshold[1])
        prob_pos = prob[idx_pos].reshape([-1, 1])
        bbox_pos = boxes[idx_pos]
        bbox_reg_pos = bbox_reg[idx_pos]
        proposals_total = np.concatenate(
            (bbox_pos, prob_pos, bbox_reg_pos), axis=1
        )
        if len(proposals_total) == 0:
            return (np.zeros((0, 5), dtype=np.float32),
                    np.zeros((0, 10), dtype=np.float32))

        proposals_total = proposals_total[
            nms(proposals_total, threshold=iou_threshold[2])
        ]
        proposals_total = bbreg(proposals_total)

        # Stage 3
        proposals_total = rerec(proposals_total)
        boxes = proposals_total[:, 0:4].astype(np.int32)
        input = cut(img, boxes, h, w, 48)
        input = (input - 127.5) * 0.0078125
        try:
            cls_prob, bbox_reg, landmark_reg = self.onet.predict(input)
        except:
            return (np.zeros((0, 5), dtype=np.float32),
                    np.zeros((0, 10), dtype=np.float32))

        prob = cls_prob[..., 1]
        idx_pos = np.where(prob >= score_threshold[2])
        prob_pos = prob[idx_pos].reshape([-1, 1])
        bbox_pos = boxes[idx_pos]
        # refine landmark
        height = bbox_pos[:, 2:3] - bbox_pos[:, 0:1] + 1
        width = bbox_pos[:, 3:4] - bbox_pos[:, 1:2] + 1
        landmark = landmark_reg[idx_pos]
        landmark[:, 0:5] = np.round(
            bbox_pos[:, 0:1] + landmark[:, 0:5] * height
        )
        landmark[:, 5:10] = np.round(
            bbox_pos[:, 1:2] + landmark[:, 5:10] * width
        )
        bbox_reg_pos = bbox_reg[idx_pos]
        proposals_total = np.concatenate(
            (bbox_pos, prob_pos, bbox_reg_pos), axis=1
        )
        if len(proposals_total) == 0:
            return (np.zeros((0, 5), dtype=np.float32),
                    np.zeros((0, 10), dtype=np.float32))

        proposals_total = bbreg(proposals_total)
        idx_keep = nms(proposals_total, threshold=iou_threshold[2],
                       type='Min')

        return proposals_total[idx_keep][:, 0:5], landmark[idx_keep]