"""Functions and classes that allow images to be differenced"""
import os
import cPickle
import numpy as np
import scipy.stats
import astropy.io.fits
from astropy.stats import sigma_clip
import Polygon
import Polygon.Utils
import multiprocessing

from matplotlib.figure import Figure
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas

class SliceableBase(object):
    """Baseclass for dealing with image slices."""

    def _apply_mask(self, image, mask):
        """Set masked pixels to NaN."""
        image[mask > 0] = np.nan
        return image

    def getGoodPix(self, fracThreshold=0.5, threshold=None):
        if threshold is None:
            threshold = fracThreshold * \
                (self.weight.max() - self.weight.min()) + self.weight.min()
        return self.weight > threshold

    def saveSegments(self, pathRoot):
        """Saves a the sliced FITS image to pathRoot_im.fits and the sliced
        weight FITS to pathRoot_wht.fits.
        """
        dirname = os.path.dirname(pathRoot)
        if os.path.exists(dirname) is False:
            os.makedirs(dirname)
        astropy.io.fits.writeto(pathRoot + "_im.fits", self.image,
                                clobber=True)
        if self.weight is not None:
            astropy.io.fits.writeto(pathRoot + "_wht.fits", self.weight,
                                    clobber=True)

class SliceableImage(SliceableBase):
    """Used by _computeDiff to efficiently slice overlapping images and compare
    their common set of good pixels.
    """
    def __init__(self, key, image, weight, mask=None):
        super(SliceableImage, self).__init__()
        self.key = key
        self.fullImage = image
        self.fullWeight = weight
        if mask is not None:
            self.fullImage = self._apply_mask(image, mask)
            self.fullWeight = self._apply_mask(weight, mask)
        self.r = None
        self.image = None
        self.weight = None

    @classmethod
    def makeFromFITS(cls, key, imagePath, weightPath, mask_path=None):
        """Create a `SliceableImage` from FITS paths"""
        print imagePath, weightPath, mask_path
        if weightPath is not None:
            w = astropy.io.fits.getdata(weightPath, 0)
        else:
            w = None
        if mask_path is not None:
            m = astropy.io.fits.getdata(mask_path, 0)
        else:
            m = None
        return cls(key, astropy.io.fits.getdata(imagePath, 0), w, mask=m)

    def _slice(self, im):
        return im[self.r[1][0]:self.r[1][1], self.r[0][0]:self.r[0][1]]

    def setRange(self, sliceRange):
        self.r = sliceRange
        self.image = self._slice(self.fullImage)
        if self.fullWeight is not None:
            self.weight = self._slice(self.fullWeight)
        else:
            self.weight = None

class SliceableFITS(SliceableBase):
    """Used for accessing slices of a large FITS image. A semi-replacement for
    using the Sliceableimage.makeFromFITS interface. This class will read the
    image slice in memory as the range is set, thus minimizing memory usage.
    """
    def __init__(self, key, imagePath, weightPath):
        super(SliceableFITS, self).__init__()
        self.key = key
        self.imagePath = imagePath
        self.weightPath = weightPath
        self.imageFITS = astropy.io.fits.open(self.imagePath)
        self.weightFITS = astropy.io.fits.open(self.weightPath)
        self.image = None
        self.weight = None
        self.r = None

    def setRange(self, sliceRange):
        self.r = sliceRange
        imageSec = self.imageFITS[0].section[self.r[1][0]:self.r[1][1], :]
        weightSec = self.weightFITS[0].section[self.r[1][0]:self.r[1][1], :]
        self.image = imageSec[:, self.r[0][0]:self.r[0][1]]
        self.weight = weightSec[:, self.r[0][0]:self.r[0][1]]

    def close(self):
        """Deallocate the FITS file references."""
        self.imageFITS.close()
        self.weightFITS.close()

class ResampledWCS(object):
    """Holds the WCS of a single resampled image frame.
    """
    def __init__(self, header):
        super(ResampledWCS, self).__init__()
        self.header = header
        self.naxis1 = self.header['NAXIS1']
        self.naxis2 = self.header['NAXIS2']
        self.crpix1 = int(self.header['crpix1']) - 1  # to be 0-based
        self.crpix2 = int(self.header['crpix2']) - 1
        self.makePolygon()

    def makePolygon(self):
        """Makes self.polygon, which is a `Polygon` instance with vertices in
        the primed space of the mosaic image.
        """
        polyPoints = self.getVertices()
        self.polygon = Polygon.Polygon(polyPoints)

    def getVertices(self):
        """docstring for getVertices"""
        print self.naxis1, self.naxis2, self.crpix1, self.crpix2
        vertsX = np.array([0, 0, self.naxis1 - 1, self.naxis1 - 1],
                          dtype=int) - self.crpix1
        vertsY = np.array([0, self.naxis2 - 1, self.naxis2 - 1, 0],
                          dtype=int) - self.crpix2
        vertsX = list(vertsX)
        vertsY = list(vertsY)
        polyPoints = zip(vertsX, vertsY)
        return polyPoints

    def getCentroid(self):
        """:return: (x,y) tuple of the frame centroid in the mosaic space."""
        x = np.array([0, 0, self.naxis1 - 1, self.naxis1 - 1], dtype=int) \
            - self.crpix1
        y = np.array([0, self.naxis2 - 1, self.naxis2 - 1, 0], dtype=int) \
            - self.crpix2
        return x.mean(), y.mean()

class ManualResampledWCS(ResampledWCS):
    """Manually create a ResampledWCS from NAXISi and CRPIXi."""
    def __init__(self, naxis1, naxis2, crpix1, crpix2):
        self.naxis1 = naxis1
        self.naxis2 = naxis2
        self.crpix1 = int(crpix1) - 1  # so that everything is zero-based
        self.crpix2 = int(crpix2) - 1
        self.makePolygon()

class Overlap(object):
    """For an overlapping pair of images, records the slice of their overlap.
    """
    def __init__(self, upperFrame, lowerFrame):
        super(Overlap, self).__init__()
        self.upperFrame = upperFrame
        self.lowerFrame = lowerFrame

        self._computeOverlap()

        if self.overlap is not None:
            self.upSlice = self._makeOverlapRangeForImageFrame(self.upperFrame)
            self.loSlice = self._makeOverlapRangeForImageFrame(self.lowerFrame)

    def _computeOverlap(self):
        """Determins the polygon of the overlap between the two frames."""
        if self.upperFrame.polygon.overlaps(self.lowerFrame.polygon):
            self.overlap = self.upperFrame.polygon & self.lowerFrame.polygon
        else:
            self.overlap = None

    def _makeOverlapRangeForImageFrame(self, frame):
        """Converts a polygon in co-add space to a range across the image's
        axes.
        """
        # Convert the overlap polygon (in the prime co-add space) to the
        # coordinates of the resampled image
        primeVerts = Polygon.Utils.pointList(self.overlap)
        xlist = []
        ylist = []
        for (xPrime, yPrime) in primeVerts:
            xlist.append(xPrime + frame.crpix1)
            ylist.append(yPrime + frame.crpix2)
        ymin = int(min(ylist))
        ymax = int(max(ylist))
        xmin = int(min(xlist))
        xmax = int(max(xlist))
        # add one to the upper limit because of the behaviour of range() func
        return (xmin, xmax + 1), (ymin, ymax + 1)

    def hasOverlap(self):
        if self.overlap is None:
            return False
        else:
            return True

    def getUpperLowerSlices(self):
        """:return: tuple (upper slice, lower slice)"""
        return self.upSlice, self.loSlice

    def getOverlapCRPIX(self, r, frame):
        """:return: tuple (CRPIX1, CRPIX2) of the overlapping image frame."""
        # The left slice of the upper image. This is the amount that the CRPIX1
        # will shift by.
        deltaX = r[0][0]
        crpix1 = frame.crpix1 - deltaX
        # Do the same with the yaxis
        deltaY = r[1][0]
        crpix2 = frame.crpix2 - deltaY
        return (crpix1, crpix2)

    def getLowerOverlapCRPIX(self):
        """docstring for getLowerOverlapCRPIX"""
        return self.getOverlapCRPIX(self.loSlice, self.lowerFrame)

    def getUpperOverlapCRPIX(self):
        """docstring for getLowerOverlapCRPIX"""
        return self.getOverlapCRPIX(self.upSlice, self.upperFrame)

    def getUpperCentreTrans(self):
        """:return: (delta x, delta y) tuple of offsets to go from the centre
        of the upper frame, to the centre of the overlap frame."""
        return self._makeCentreTrans(self.upperFrame, self.upSlice)

    def getLowerCentreTrans(self):
        """:return: (delta x, delta y) tuple of offsets to go from the centre
        of the upper frame, to the centre of the overlap frame."""
        return self._makeCentreTrans(self.lowerFrame, self.loSlice)

    def _makeCentreTrans(self, frame, slce):
        oXMin, oXMax = slce[0]
        oYMin, oYMax = slce[1]
        oWidth = oXMax - oXMin
        oHeight = oYMax - oYMin

        fWidth = frame.naxis1
        fHeight = frame.naxis2

        xDelta = oXMin + oWidth / 2. - fWidth / 2.
        yDelta = oYMin + oHeight / 2. - fHeight / 2.
        return xDelta, yDelta

class OverlapDB(object):
    """Container for finding and holding all overlaps between a set of
    ResampledWCS.
    """
    def __init__(self, resampledWCSs):
        super(OverlapDB, self).__init__()
        # frame/field dictionary of ResampledWCS instances
        self.resampledWCSs = resampledWCSs
        self._findOverlaps()

    def __getitem__(self, targetField):
        """:return: a list of all fields that have overlaps with `targetField`.
        """
        coupledFields = []
        for overlapKey in self.overlaps.keys():
            if targetField in overlapKey:
                if targetField == overlapKey[0]:
                    coupledFields.append(overlapKey[1])
                else:
                    coupledFields.append(overlapKey[0])
        return list(set(coupledFields))

    def _findOverlaps(self):
        """docstring for findOverlaps"""
        self.overlaps = {}
        frameIDs = self.resampledWCSs.keys()
        for i, frameID1 in enumerate(frameIDs):
            for frameID2 in frameIDs[i + 1:]:
                overlap = Overlap(self.resampledWCSs[frameID1],
                                  self.resampledWCSs[frameID2])
                if overlap.hasOverlap():
                    overlapKey = (frameID1, frameID2)
                    self.overlaps[overlapKey] = overlap

    def iteritems(self):
        """Delegate iteritems to the overlaps dictionary"""
        return self.overlaps.iteritems()

class Couplings(object):
    """Computes and stores image differences for overlapping fields."""
    def __init__(self):
        super(Couplings, self).__init__()
        self.fields = {}
        self.fieldDiffs = {}
        self.fieldDiffSigmas = {}
        self.fieldDiffAreas = {}
        self.fieldLevels = {}
        self.coverage_fractions = {}
        self.diff_paths = {}

    def __getitem__(self, fieldname):
        """:return: list of fields that are coupled to the named field."""
        return self.overlapDB[fieldname]

    def get_doc(self):
        """Makes a dictionary document for saving the couplings in MongoDB.

        The field tuples are replaced with a single string key using '*'
        as the delimter.
        """
        diffs = {}
        sigmas = {}
        areas = {}
        levels = {}
        coverage_fractions = {}
        diff_paths = {}
        for pair, diff in self.fieldDiffs.iteritems():
            diffs["*".join(pair)] = float(diff)
        for pair, sigma in self.fieldDiffSigmas.iteritems():
            sigmas["*".join(pair)] = float(sigma)
        for pair, area in self.fieldDiffAreas.iteritems():
            areas["*".join(pair)] = float(area)
        for pair, level in self.fieldLevels.iteritems():
            levels["*".join(pair)] = float(level)
        for pair, f in self.coverage_fractions.iteritems():
            coverage_fractions["*".join(pair)] = float(f)
        for pair, p in self.diff_paths.iteritems():
            diff_paths["*".join(pair)] = p
        doc = {"fields": self.fields,
               "diffs": diffs,
               "sigmas": sigmas,
               "areas": areas,
               "levels": levels,
               "coverage_fractions": coverage_fractions,
               "diff_paths": diff_paths}
        return doc

    @classmethod
    def load_doc(cls, doc):
        """Make a Couplings instance from a MongoDB document.

        The mongoDB document simulates tuples of field pairs as a string with
        '*' as the delimeter.
        """
        instance = cls()
        instance.fields = doc['fields']
        instance.fieldDiffs = {}
        instance.fieldDiffSigmas = {}
        instance.fieldDiffAreas = {}
        instance.fieldLevels = {}
        instance.coverage_fractions = {}
        instance.diff_paths = {}
        for fieldPair, subDoc in doc['diffs'].iteritems():
            fieldTuple = tuple(fieldPair.split("*"))
            instance.fieldDiffs[fieldTuple] = subDoc
        for fieldPair, subDoc in doc['sigmas'].iteritems():
            fieldTuple = tuple(fieldPair.split("*"))
            instance.fieldDiffSigmas[fieldTuple] = subDoc
        for fieldPair, subDoc in doc['areas'].iteritems():
            fieldTuple = tuple(fieldPair.split("*"))
            instance.fieldDiffAreas[fieldTuple] = subDoc
        for fieldPair, subDoc in doc['levels'].iteritems():
            fieldTuple = tuple(fieldPair.split("*"))
            instance.fieldLevels[fieldTuple] = subDoc
        for fieldPair, subDoc in doc['coverage_fractions'].iteritems():
            fieldTuple = tuple(fieldPair.split("*"))
            instance.coverage_fractions[fieldTuple] = subDoc
        for fieldPair, subDoc in doc['diff_paths'].iteritems():
            fieldTuple = tuple(fieldPair.split("*"))
            instance.diff_paths[fieldTuple] = subDoc
        return instance

    def save(self, path):
        """Pickle the couplings object to the given `path`."""
        dirname = os.path.dirname(path)
        if os.path.exists(dirname) is False:
            os.makedirs(dirname)
        if os.path.exists(path):
            os.remove(path)

        f = open(path, 'w')
        p = cPickle.Pickler(f)
        p.dump(self.fieldDiffs)
        p.dump(self.fieldDiffSigmas)
        p.dump(self.fieldDiffAreas)
        p.dump(self.fieldLevels)
        p.dump(self.coverage_fractions)
        p.dump(self.diff_paths)
        f.close()

    @classmethod
    def load(cls, path, fields):
        """Load a Couplings object from a pickle at the given `path`. Fields
        is a dictionary of fieldname: field object. Only couplings involving
        these particular fields will be used.
        """
        f = open(path, 'r')
        p = cPickle.Unpickler(f)
        instance = cls()
        instance.fields = fields
        instance.fieldDiffs = p.load()
        instance.fieldDiffSigmas = p.load()
        instance.fieldDiffAreas = p.load()
        instance.fieldLevels = p.load()
        instance.coverage_fractions = p.load()
        instance.diff_paths = p.load()
        f.close()

        instance._trim_diffs_to_fields()
        return instance

    def _trim_diffs_to_fields(self):
        """When loading a couplings object from a pickle, you may want to
        use fewer field objects than went into the original couplings DB.
        """
        fieldNames = self.fields.keys()
        pairKeys = self.fieldDiffs.keys()
        for pairKey in pairKeys:
            field1, field2 = pairKey
            if (field1 not in fieldNames) or (field2 not in fieldNames):
                del self.fieldDiffs[pairKey]
                del self.fieldDiffSigmas[pairKey]
                del self.fieldDiffAreas[pairKey]
                del self.fieldLevels[pairKey]
                del self.coverage_fractions[pairKey]
                del self.diff_paths[pairKey]

    def get_overlap_db(self):
        return self.overlapDB

    def add_field(self, name, imagePath, weightPath, mask_path=None):
        """Adds a field object, which is any object that responds to the field
        specification...
        """
        self.fields[name] = {'image_path': imagePath,
                             'weight_path': weightPath,
                             'mask_path': mask_path}

    def make(self, diffDir, mp=True, plotDir=None):
        """Compute coupled image differences, using `diffDir` as a working
        directory for the image differencing.

        Parameters
        ----------
        diffDir : str
            Directory where difference images are saved.
        mp : bool
            If True, then difference images are computed in parallel.
        """
        self._find_overlaps()
        self._compute_overlap_diffs(diffDir, mp=mp, plotDir=plotDir)

    def _find_overlaps(self):
        """Given the fields, detects and records all overlaps between fields.
        """
        fieldFootprints = {}
        for fieldname, field in self.fields.iteritems():
            fieldFootprints[fieldname] = self._make_imageframe(
                field['image_path'])
        self.overlapDB = OverlapDB(fieldFootprints)

    def _make_imageframe(self, imagePath):
        """:return: the ResampledWCS corresponding to image (assumed in ext 0).
        """
        resampledWCS = ResampledWCS(astropy.io.fits.getheader(imagePath, 0))
        return resampledWCS

    def _compute_overlap_diffs(self, diffImageDir, mp=True, plotDir=None):
        """Compute the scalar offsets between fields in each overlap.

        This is done with multiprocessing, so that overlaps are computed in
        parallel.

        :param diffImageDir: can optionally be set to a directory name so
            that the computed difference images will be saved there.
        """
        if plotDir is None:
            plotDir = diffImageDir
        if os.path.exists(diffImageDir) is False:
            os.makedirs(diffImageDir)
        if os.path.exists(plotDir) is False:
            os.makedirs(plotDir)

        args = []
        for overlapKey, overlap in self.overlapDB.iteritems():
            field1, field2 = overlapKey
            field1Path = self.fields[field1]['image_path']
            field1WeightPath = self.fields[field1]['weight_path']
            field1MaskPath = self.fields[field1]['mask_path']
            field2Path = self.fields[field2]['image_path']
            field2WeightPath = self.fields[field2]['weight_path']
            field2MaskPath = self.fields[field2]['mask_path']
            arg = (field1, field1Path, field1WeightPath, field1MaskPath,
                   field2, field2Path, field2WeightPath, field2MaskPath,
                   overlap, diffImageDir, plotDir, 5)
            args.append(arg)
        if mp:
            pool = multiprocessing.Pool(processes=multiprocessing.cpu_count())
            results = pool.map(_computeDiff, args)
            pool.close()
            pool.join()
        else:
            results = map(_computeDiff, args)
        diffs = {}
        diffSigmas = {}
        diffAreas = {}
        meanLevels = {}
        coverage_fractions = {}
        diff_paths = {}
        for result in results:
            field1, field2, offsetData = result
            pairKey = (field1, field2)
            if offsetData is not None:
                diffs[pairKey] = offsetData['mean']
                diffSigmas[pairKey] = offsetData['sigma']
                diffAreas[pairKey] = offsetData['area']
                meanLevels[pairKey] = offsetData['level']
                coverage_fractions[pairKey] = offsetData['coverage_frac']
                diff_paths[pairKey] = offsetData['diff_path']
            else:
                # self.overlaps.rejectOverlap(pairKey)
                pass
        self.fieldDiffs = diffs
        self.fieldDiffSigmas = diffSigmas
        self.fieldDiffAreas = diffAreas
        self.fieldLevels = meanLevels
        self.coverage_fractions = coverage_fractions
        self.diff_paths = diff_paths

    def get_field_diffs(self, omittedFields=[], replicatedFields=[]):
        """Returns a dictionary of field differences, and a list of fields
        represented.

        Parameters
        ----------
        omittedFields : list
            A list of fields whose couplings should be omitted
            from the list of field dfferences.
        replicatedFields : list
            A list of fields whose couplings should be repeated
            (ie, given more weight.) A field can be replicated an arbitrary
            number of times. This is useful for jackknife/bootstrap analysis.
        """
        # Split the fieldDiffs dictionary into lists of keys and values
        # (this is done because replication will require duplicated keys)
        pairNames = []
        diffs = []
        diffSigmas = {}
        for pairName, diff in self.fieldDiffs.iteritems():
            pairNames.append(pairName)
            diffs.append(diff)
            diffSigmas[pairName] = self.fieldDiffSigmas[pairName]
        print "There are %i pairs initially" % len(pairNames)
        pairNames, diffs = self._getFieldDiffsAfterReplication(
            pairNames, diffs, replicatedFields)
        print "There are %i pairs after replication" % len(pairNames)
        pairNames, diffs = self._getFieldDiffsAfterOmission(
            pairNames, diffs, omittedFields)
        print "There are %i pairs after omission" % len(pairNames)
        print "Should have omitted %i fields" % len(omittedFields)
        # Now get a set of the represented fields
        allFields = []
        for pairKey in pairNames:
            allFields.append(pairKey[0])
            allFields.append(pairKey[1])
        allFields = list(set(allFields))
        return pairNames, diffs, diffSigmas, allFields

    def _getFieldDiffsAfterReplication(self, pairNames, diffs,
                                       replicatedFields):
        """Duplicates items in pairNames and diffs for each match item in
        replicatedFields.
        """
        extraPairNames = []
        extraDiffs = []
        for replicatedField in replicatedFields:
            for i, pairName in enumerate(pairNames):
                if replicatedField in pairName:
                    # duplicate this coupling
                    extraPairNames.append(pairNames[i])
                    extraDiffs.append(diffs[i])
        pairNames += extraPairNames
        diffs += extraDiffs
        return pairNames, diffs

    def _getFieldDiffsAfterOmission(self, pairNames, diffs, omittedFields):
        """Omits an coupling that includes the omitted fields."""
        retainedPairNames = []
        retainedDiffs = []
        for i, pairName in enumerate(pairNames):
            keep = True
            for field in pairName:
                if field in omittedFields:
                    keep = False
            if keep is True:
                retainedPairNames.append(pairNames[i])
                retainedDiffs.append(diffs[i])
        return retainedPairNames, retainedDiffs

def _computeDiff(arg):
    """Worker: Computes the DC offset of frame-coadd"""
    upperKey, upperPath, upperWeightPath, upperMaskKey, lowerKey, lowerPath, \
        lowerWeightPath, lowerMaskKey, \
        overlap, diffDir, diffPlotDir, nsigma = arg
    upper = SliceableImage.makeFromFITS(
        upperKey, upperPath, upperWeightPath, mask_path=upperMaskKey)
    upper.setRange(overlap.upSlice)
    lower = SliceableImage.makeFromFITS(
        lowerKey, lowerPath, lowerWeightPath, mask_path=lowerMaskKey)
    lower.setRange(overlap.loSlice)
    goodPix = np.where((upper.weight > 0.) & (lower.weight > 0.))
    nPixels = len(goodPix[0])
    print "diff pair", lowerKey, upperKey, nPixels
    if nPixels > 10:
        # Offset via difference image
        diff_pixels = upper.image[goodPix] - lower.image[goodPix]
        # Choose iters=None so clipping until convergence.
        clipped = sigma_clip(diff_pixels, sig=nsigma, iters=1,
                             varfunc=np.nanvar)
        median = np.nanmedian(clipped[~clipped.mask])
        sigma = np.nanstd(clipped[~clipped.mask])
        n_clipped_pixels = len(clipped[~clipped.mask])
        cov_frac = float(n_clipped_pixels) / len(diff_pixels.ravel())
        print "Offset %.2e +/- %.2e" % (median, sigma)
        offsetData = {"mean": median,
                      "sigma": sigma,
                      "area": n_clipped_pixels,
                      "level": median,
                      "coverage_frac": cov_frac}

        # Save the difference image, if possible
        if diffDir is not None:
            image = upper.image - lower.image
            bad_pix = np.where((upper.weight == 0.) | (lower.weight == 0.))
            image[bad_pix] = np.nan
            lower_lim = median - nsigma * sigma
            upper_lim = median + nsigma * sigma
            image[image < lower_lim] = np.nan
            image[image > upper_lim] = np.nan
            path = os.path.join(diffDir, "%s*%s.fits" % (upperKey, lowerKey))
            astropy.io.fits.writeto(path, image, clobber=True)
            offsetData['diff_path'] = path
    else:
        offsetData = None
        path = None
    return upperKey, lowerKey, offsetData

class CoupledPlanes(object):
    """A replacement for the `Couplings` class in the context of slope
    analysis.
    """
    def __init__(self):
        super(CoupledPlanes, self).__init__()
        self.fields = {}
        self.diffPlanes = {}
        self.diffAreas = {}
        self.diffSigmas = {}
        # mean DC level in the overlap
        self.fieldLevels = {}
        # displacement in (x,y) pixels from the centre of the upper image to
        # the overlap centre
        self.upperOverlapTrans = {}
        # same, for the lower image.
        self.lowerOverlapTrans = {}
        self.overlapShape = {}  # (xsize, ysize) of the overlap area

    def __getitem__(self, fieldname):
        """:return: a list of fields that are coupled to the same named field.
        """
        return self.overlapsDB[fieldname]

    def get_doc(self):
        """Makes a dictionary document for saving the couplings in MongoDB.

        The field tuples are replaced with a single string key using '*'
        as the delimter.
        """
        diffs = {}
        sigmas = {}
        areas = {}
        levels = {}
        upperTrans = {}
        lowerTrans = {}
        overlapShape = {}
        for pair, diff in self.diffPlanes.iteritems():
            diffs["*".join(pair)] = diff
        for pair, sigma in self.diffSigmas.iteritems():
            sigmas["*".join(pair)] = sigma
        for pair, area in self.diffAreas.iteritems():
            areas["*".join(pair)] = area
        for pair, level in self.fieldLevels.iteritems():
            levels["*".join(pair)] = level
        for pair, trans in self.upperOverlapTrans.iteritems():
            upperTrans["*".join(pair)] = trans
        for pair, trans in self.lowerOverlapTrans.iteritems():
            lowerTrans["*".join(pair)] = trans
        for pair, shape in self.overlapShape.iteritems():
            overlapShape["*".join(pair)] = shape

        doc = {"fields": self.fields,
               "diffs": diffs,
               "sigmas": sigmas,
               "areas": areas,
               "levels": levels,
               "upper_trans": upperTrans,
               "lower_trans": lowerTrans,
               "shape": overlapShape}
        return doc

    @classmethod
    def load_doc(cls, doc):
        """Make a CoupledPlanes instance from a MongoDB document.

        The mongoDB document simulates tuples of field pairs as a string with
        '*' as the delimeter.
        """
        instance = cls()
        print "loading CoupledPlanes doc", doc
        instance.fields = doc['fields']
        instance.diffPlanes = {}
        instance.diffSigmas = {}
        instance.diffAreas = {}
        instance.fieldLevels = {}
        instance.upperOverlapTrans = {}
        instance.lowerOverlapTrans = {}
        instance.overlapShape = {}
        for fieldPair, subDoc in doc['diffs'].iteritems():
            fieldTuple = tuple(fieldPair.split("*"))
            instance.diffPlanes[fieldTuple] = subDoc
        for fieldPair, subDoc in doc['sigmas'].iteritems():
            fieldTuple = tuple(fieldPair.split("*"))
            instance.diffSigmas[fieldTuple] = subDoc
        for fieldPair, subDoc in doc['areas'].iteritems():
            fieldTuple = tuple(fieldPair.split("*"))
            instance.diffAreas[fieldTuple] = subDoc
        for fieldPair, subDoc in doc['levels'].iteritems():
            fieldTuple = tuple(fieldPair.split("*"))
            instance.fieldLevels[fieldTuple] = subDoc
        for fieldPair, subDoc in doc['upper_trans'].iteritems():
            fieldTuple = tuple(fieldPair.split("*"))
            instance.upperOverlapTrans[fieldTuple] = subDoc
        for fieldPair, subDoc in doc['lower_trans'].iteritems():
            fieldTuple = tuple(fieldPair.split("*"))
            instance.lowerOverlapTrans[fieldTuple] = subDoc
        for fieldPair, subDoc in doc['shape'].iteritems():
            fieldTuple = tuple(fieldPair.split("*"))
            instance.overlapShape[fieldTuple] = subDoc

        return instance

    def add_field(self, name, imagePath, weightPath):
        """Adds a field object, which is any object that responds to the field
        specification..."""
        self.fields[name] = {'image_path': imagePath,
                             'weight_path': weightPath}

    def make(self, diffDir, nThreads=None):
        """Compute coupled image differences, using `diffDir` as a working
        directory for the image differencing.
        """
        self._find_overlaps()
        self._compute_overlap_diffs(diffDir, nProcesses=nThreads)

    def _find_overlaps(self):
        """Given the fields, detects and records all overlaps between fields.
        """
        fieldFootprints = {}
        for fieldname, field in self.fields.iteritems():
            fieldFootprints[fieldname] = self._make_imageframe(
                field['image_path'])
        self.overlapDB = OverlapDB(fieldFootprints)

    def _make_imageframe(self, imagePath):
        """:return: the ResampledWCS corresponding to image (assumed in ext 0).
        """
        resampledWCS = ResampledWCS(astropy.io.fits.getheader(imagePath, 0))
        return resampledWCS

    def get_overlap_db(self):
        return self.overlapDB

    def _compute_overlap_diffs(self, diffImageDir, nProcesses=None):
        """Compute the difference planes between fields in each overlap.

        Done with multiprocessing, so that overlaps are computed in parallel.

        :param diffImageDir: can be set to a directory name so that the
            computed difference images will be saved there.
        """
        if diffImageDir is not None:
            if os.path.exists(diffImageDir) is False:
                os.makedirs(diffImageDir)

        args = []
        for overlapKey, overlap in self.overlapDB.iteritems():
            field1, field2 = overlapKey
            field1Path = self.fields[field1]['image_path']
            field1WeightPath = self.fields[field1]['weight_path']
            field2Path = self.fields[field2]['image_path']
            field2WeightPath = self.fields[field2]['weight_path']
            arg = (field1, field1Path, field1WeightPath, field2, field2Path,
                   field2WeightPath, overlap, diffImageDir)
            args.append(arg)

        if nProcesses is None:
            nProcesses = multiprocessing.cpu_count()
        pool = multiprocessing.Pool(processes=nProcesses)
        results = pool.map(_computeDiffPlane, args)
        # results = map(_computeDiffPlane, args)

        self.diffPlanes = {}
        self.diffAreas = {}
        self.diffSigmas = {}
        self.fieldLevels = {}
        self.upperOverlapTrans = {}
        self.lowerOverlapTrans = {}
        self.overlapShape = {}
        for result in results:
            field1, field2, diffData = result
            pairKey = (field1, field2)
            if diffData is not None:
                self.diffPlanes[pairKey] = diffData['plane']
                self.diffSigmas[pairKey] = diffData['sigma']
                self.diffAreas[pairKey] = diffData['area']
                self.fieldLevels[pairKey] = diffData['level']
                self.upperOverlapTrans[pairKey] = diffData['upper_trans']
                self.lowerOverlapTrans[pairKey] = diffData['lower_trans']
                self.overlapShape[pairKey] = diffData['shape']

def _computeDiffPlane(arg):
    """Worker for computing the difference plane between two coupled images."""
    upperKey, upperPath, upperWeightPath, lowerKey, lowerPath, \
        lowerWeightPath, overlap, diffDir = arg
    upper = SliceableImage.makeFromFITS(upperKey, upperPath, upperWeightPath)
    upper.setRange(overlap.upSlice)
    lower = SliceableImage.makeFromFITS(lowerKey, lowerPath, lowerWeightPath)
    lower.setRange(overlap.loSlice)
    goodPix = np.where((upper.weight > 0.) & (lower.weight > 0.))
    nPixels = len(goodPix[0])

    if nPixels > 10:
        # Iterate several times until convergence is reached; we're really
        # iterating to build a good description of what the good pixels are
        # Decompose the image into a list of pixels, and their indices
        diffImage = upper.image - lower.image
        shape = diffImage.shape
        ysize, xsize = shape
        y0 = int(ysize / 2.)
        x0 = int(xsize / 2.)

        xCoords = []
        yCoords = []
        xIndices = []
        yIndices = []
        # flagPix = []
        for i in xrange(xsize):
            x = i - x0
            for j in xrange(ysize):
                y = j - y0
                xCoords.append(x)
                yCoords.append(y)
                xIndices.append(i)
                yIndices.append(j)
        xCoords = np.array(xCoords)
        yCoords = np.array(yCoords)
        xIndices = np.array(xIndices, dtype=int)
        yIndices = np.array(yIndices, dtype=int)

        converged = False
        planeFit, newGoodPix = _fitPlane(diffImage, xCoords, yCoords,
                                         xIndices, yIndices, goodPix)
        newGoodPix = _combineGoodPixMasks(diffImage.shape, goodPix, newGoodPix)
        nIter = 0
        while converged is False:
            newPlaneFit, newGoodPix = _fitPlane(diffImage, xCoords, yCoords,
                                                xIndices, yIndices, newGoodPix)

            # Need to combine newGoodPix, which is based purely on sigma clip
            # with the original goodPix. The surface brightness can often be
            # less than the difference flux. Thus, a sigma-clipped newGoodPix
            # will often include non-difference pixels
            newGoodPix = _combineGoodPixMasks(diffImage.shape, goodPix,
                                              newGoodPix)

            planePerDiff = (planeFit - newPlaneFit) / planeFit
            planePerDiffSq = np.sum(planePerDiff ** 2.)
            if planePerDiffSq < 1e-8:
                converged = True
                print "Converged on iter %i" % nIter
            planeFit = newPlaneFit
            nIter += 1
            if nIter > 10:
                break
                print "Iteration timeout"

        print "Plane", planeFit

        # Get metadata of the couplings
        area = len(goodPix[0])
        meanLevel = np.median((upper.image[newGoodPix]
                               + lower.image[newGoodPix]) / 2.)
        # standard deviation of the diffence image to plane
        args = diffImage, xCoords, yCoords, xIndices, yIndices, newGoodPix
        sigma = np.sqrt(_planeObjFunc(planeFit, *args) ** 2.).std()

        # The translation to go from the centre of the upper frame to
        # the centre of the overlap
        upperTrans = overlap.getUpperCentreTrans()
        lowerTrans = overlap.getLowerCentreTrans()

        offsetData = {"plane": tuple(planeFit), "sigma": sigma, "area": area,
                      "level": meanLevel, "upper_trans": upperTrans,
                      "lower_trans": lowerTrans, "shape": (xsize, ysize)}

        # Save the difference image
        outputImage = np.nan * np.zeros(diffImage.shape)
        outputImage[newGoodPix] = diffImage[newGoodPix]
        astropy.io.fits.writeto(os.path.join(diffDir,
                                "%s_%s.fits" % (upperKey, lowerKey)),
                                outputImage, clobber=True)
    else:
        offsetData = None
    return upperKey, lowerKey, offsetData

def _fitPlane(diffImage, xCoords, yCoords, xIndices, yIndices, goodPix):
    # Fit the image using least squares
    args = diffImage, xCoords, yCoords, xIndices, yIndices, goodPix
    p0 = [0., 0., 0.]  # mx, my, offset
    p, success = scipy.optimize.leastsq(_planeObjFunc, p0, args=args)

    # Build new mask from sigma clipping the difference of model to real image
    mx, my, c = p
    modelImage = np.zeros(diffImage.shape)
    modelImage[yIndices, xIndices] = xCoords * mx + yCoords * my + c
    delta = diffImage - modelImage
    sigma = delta.std()
    goodPix = np.where(delta ** 2. < sigma ** 2.)  # clip 1-sigma deviations
    return p, goodPix

def _planeObjFunc(p, *args):
    realImage, xCoords, yCoords, xIndices, yIndices, goodPix = args

    mx, my, c = p
    modelImage = np.zeros(realImage.shape)
    modelImage[yIndices, xIndices] = xCoords * mx + yCoords * my + c

    delta = realImage - modelImage
    return np.ravel(delta[goodPix])

def _combineGoodPixMasks(shape, goodPix, newGoodPix):
    goodImage = np.zeros(shape, dtype=int)
    goodImage[goodPix] = 1

    newGoodImage = np.zeros(shape, dtype=int)
    newGoodImage[newGoodPix] = 1

    # Pixels need to be good in *both* masks
    combinedImage = goodImage & newGoodImage
    combinedGoodPix = np.where(combinedImage == 1)
    return combinedGoodPix

def histogramFit(pixels, initMean, initSigma):
    """Fits the offset and its error by fitting a guassian to the histogram."""
    nPixels = pixels.shape[0]
    nBins = int(nPixels / 1000)
    hist, bins = np.histogram(pixels, bins=nBins)
    binCentres = _makeBinMiddles(bins)
    fitSigma, fitPeak, fitMean = _fitBinnedGaussian(binCentres, hist, initMean,
                                                    initSigma)
    fitSigma = np.sqrt(fitSigma ** 2.)
    return fitPeak, fitSigma, fitMean

def _makeBinMiddles(bins):
    middles = (bins[0:-1] + bins[1:]) / 2.
    return middles

unit_gauss = lambda x, mean, sigma: np.exp(-(mean - x) ** 2
                                           / (2. * sigma ** 2.))

def _fitBinnedGaussian(bins, counts, initMean, initSigma):
    """Fits a Gaussian directly to a histogram, not the moments of the
    original observations.
    """
    initialGuess = np.array([initSigma, counts.max(), initMean],
                            dtype=float)
    fit = scipy.optimize.fmin(_gaussianObjective, initialGuess,
                              args=(bins, counts))
    return fit

def _gaussianObjective(fitParams, bins, counts):
    sigma = fitParams[0]
    peak = fitParams[1]
    mean = fitParams[2]

    fittedCounts = peak * unit_gauss(bins, mean, sigma)
    diff = np.sum((fittedCounts - counts) ** 2.)
    return diff

def _diffHist(pixels, mean, clippedMedian, sigma, upperKey, lowerKey,
              plotPath):
    fig = Figure(figsize=(6, 6))
    canvas = FigureCanvas(fig)
    fig.subplots_adjust(left=0.15, bottom=0.13, wspace=0.25, hspace=0.25,
                        right=0.95)
    ax = fig.add_subplot(111)
    nPixels = len(pixels)
    # Plot histogram
    nCounts, bins, patches = ax.hist(pixels,
                                     bins=int(nPixels / 1000),
                                     histtype='stepfilled', fc='0.5',
                                     log=True, normed=True)
    # Plot estimates of the image difference
    ax.axvline(mean, ls="-", c='k', label="mean")
    ax.axvline(clippedMedian, ls='--', c='k', label="clipped median")
    # ax.axvline(fittedMean, ls='--', c='r', label="fitted mean")
    ax.legend()
    ax.text(0.05, 0.95,
            r"Clipped $\Delta= %.2e \pm %.2e$" % (clippedMedian, sigma),
            ha="left", va="top", transform=ax.transAxes)
    ax.set_xlabel("Image Difference (counts)")
    upperKey = upperKey.replace("_", "\_")
    lowerKey = lowerKey.replace("_", "\_")
    ax.set_title("%s - %s" % (upperKey, lowerKey))
    canvas.print_figure(plotPath)

class CouplingGraph(object):
    """Allows for couplings to be treated as a directed graph, with trusted
    seeds and descendents.

    An example use for this may be that one wants to build an inter-field sky
    offset solution using only a few fields in the core of the galaxy,
    where the S/N is highest. One would start with a seed as the central field,
    and follow the network until a desired number of fields are added.

    This code has some similaries to the Couplings/PropagationSolutionNetwork
    classes in DCPropagator.py. The difference is that this class uses
    the `Couplings` object to simply record all overlaps. This class uses that
    in composition, and offers the ability to then iterate through the network,
    from the seed, generation by generation.
    """
    def __init__(self):
        super(CouplingGraph, self).__init__()
        self.pastFieldSet = set([])

    def setOverlaps(self, overlapDB):
        """Set a `OverlapDB` object that defines all overlaps between fields.
        """
        self.overlapDB = overlapDB

    def setSeed(self, seedname):
        """Sets a single field as the seed."""
        self.seeds = [seedname]
        self.pastFieldSet = set(self.seeds)

    def setSeeds(self, seednames):
        """Sets multiple fields as the trusted seed."""
        self.seeds = seednames
        self.pastFieldSet = set(self.seeds)

    def __iter__(self):
        return self

    def next(self):
        """Iterates through the network, starting from the seed, generation by
        generation.
        """
        # Get the set of fields that are adjacent to fields which are
        # considered 'parents' in the network
        neighboursOfPast = []
        for pastField in self.pastFieldSet:
            neighboursOfPast += self.overlapDB[pastField]
        descendents = set(neighboursOfPast)
        newDescendents = descendents - self.pastFieldSet  # remove prior gens
        if len(newDescendents) == 0:
            raise StopIteration
        self.pastFieldSet = self.pastFieldSet | newDescendents  # update state
        return descendents

import numpy,import scipy,import astropy,from astropy


from werkzeug.utils import secure_filename
from flask import Flask, flash, request, redirect, render_template, url_for
import pdfplumber
import requests
import urllib.request
import os
from flask.helpers import send_file
import extract_text_from_pdf
from extract_text_from_pdf import get_text
import tabnanny
from summarise import top_sentence
import keywords_model as kmd
import spacy
from collections import Counter
from string import punctuation
import joblib
import numpy as np
from auto_eval import get_eval
from get_similarity import get_similarity_score

nlp = spacy.load("en_core_web_sm")
dropout_model = joblib.load('./models/try1.pkl')
counsel_model = joblib.load('./models/xgbmodel_cp.dat')

app = Flask(__name__)
app.secret_key = "secret key"
app.config['UPLOAD_FOLDER'] = './uploads'

@app.route('/summary', methods=['POST'])
def summary():
    if 'file' not in request.files:
        return {"data": "none"}

    file = request.files['file']

    ext_txt = get_text(file)
    summary = top_sentence(ext_txt, 2)
    vocab_dict, arr = kmd.textProcessing(ext_txt)
    tf = kmd.computeTF(vocab_dict, arr)
    idf = kmd.computeIDF([vocab_dict])
    tfidf = kmd.computeTfidf(tf, idf)
    keyword_ext = tfidf

    return {"value": ext_txt, "summary_text": summary, "keywords": keyword_ext}

@app.route('/plagarism', methods=['POST'])
def plagarism():
    if 'file1' not in request.files or 'file2' not in request.files:
        flash('No file part')
        return redirect(request.url)

    file = request.files['file1']

    file1 = request.files['file2']

    act_txt = get_text(file)
    que_txt = get_text(file1)
    sim_score = get_similarity_score(act_txt, que_txt)

    return {"act_text": act_txt, "que_text": que_txt, "similar": sim_score}

@app.route('/dropout', methods=['POST'])
def dropout():
    # Input retrieve
    data = request.form

    gender = int(data['gender'])
    cast = int(data['cast'])
    disabled = int(data['disabled'])
    att = int(data['att'])
    marks = list(map(int, data['marks'].split()))

    input_val = [gender, cast, disabled, att] + marks
    prediction = dropout_model.predict(np.array(input_val).reshape(1,-1))
    return {'prediction' : int(prediction[0])}

@app.route('/counsel', methods=['POST'])
def counsel():
    data = request.json
    input_val = []
    vals = [
        'sslc',
        'hsc',
        'cgpa',
        'school_type',
        'no_of_miniprojects',
        'no_of_projects',
        'coresub_skill',
        'aptitude_skill',
        'problemsolving_skill',
        'programming_skill',
        'abstractthink_skill',
        'design_skill',
        'first_computer',
        'first_program',
        'lab_programs',
        'ds_coding',
        'technology_used',
        'sympos_attend',
        'sympos_won',
        'extracurricular',
        'learning_style',
        'college_bench',
        'clg_teachers_know',
        'college_performence',
        'college_skills']
    for s in vals:
        input_val.append(int(data[s]))
    prd = counsel_model.predict([input_val])[0]
    role=[
        "Business Analyst",
        "Data Analyst",
        "Software Developer",
        "Software Tester",
        "Technical Support",
        "Technical Writer",
        "UI/UX Designer",
        "Web Developer"]
    return {'role' : role[prd]}

@app.route('/evaluate', methods=['POST'])
def evaluate():
    if 'file1' not in request.files and 'file2' not in request.files:
        flash('No file part')
        return redirect(request.url)

    file = request.files['file1']
    file1 = request.files['file2']

    act_txt = get_text(file)
    que_txt = get_text(file1)
    score = get_eval(act_txt, que_txt)

    return {'act_text': act_txt, 'que_text': que_txt, 'percentage': score[0], 'marks': score[1]}

if __name__ == "__main__":
    app.run(debug='on')


from pathlib import Path

import numpy as np

import projection_icosahedron as pro_ico
import spherical_coordinates as sc
from utility import polygon
from utility import image_io

from logger import Logger

log = Logger(__name__)
log.logger.propagate = False

class PlotFigure():
    """
    This class is used to draw figure for paper.
    """

    def __init__(self):
        self.icosahedron_radius = 1.0
        self.output_dir = None

    def draw_ico_tangent_planes_texture_stitch(self, root_dir, subimage_filename_list, texture_filename):
        """ Stitch all 20 subimage rgb image (W * H) to a single long image ( 8W * H).

        :return:
        :rtype: numpy
        """
        subimage_data_list = []
        subimage_width = -1
        subimage_height = -1
        subimage_channel = None
        for subimage_filename in subimage_filename_list:
            subimage_data = image_io.image_read(root_dir + subimage_filename)
            subimage_data_list.append(subimage_data)

            subimage_width = subimage_data.shape[1]
            subimage_height = subimage_data.shape[0]
            subimage_channel = subimage_data.shape[2]

            # if texture_height < 0:
            #     texture_height = subimage_height
            # elif texture_height > 0 and texture_height != subimage_height:
            #     log.error("the subimage size is different!")

        if subimage_channel != 3:
            log.warn("Remove the subimage alpha channel.")

        texture_height = subimage_height * 4
        texture_width = subimage_width * 5

        texture_data = np.zeros((texture_height, texture_width, 3), dtype=np.uint8)
        texture_subimage_row_start = 0

        # for idx in range(len(subimage_filename_list)):
        #     texture_subimage_end = texture_subimage_start + subimage_data_list[idx].shape[1]

        #     texture_subimage_start = texture_subimage_end
        for row_idx in range(4):
            texture_subimage_row_end = texture_subimage_row_start + subimage_height
            texture_subimage_col_start = 0
            for col_idx in range(5):
                texture_subimage_col_end = texture_subimage_col_start + subimage_width
                texture_data[texture_subimage_row_start:texture_subimage_row_end, texture_subimage_col_start: texture_subimage_col_end, :] \
                    = subimage_data_list[row_idx * 5 + col_idx][:, :, 0:3]

                texture_subimage_col_start = texture_subimage_col_end
            texture_subimage_row_start = texture_subimage_row_end

        # return texture_data
        image_io.image_save(texture_data, root_dir + texture_filename)
        # im.save('uncompressed.tga', compression=None)

    def draw_ico_tangent_planes(self, radius, padding_size, obj_file_path, subimage_shift_ratio=0.0,         obj_file_texture_enable=False, texture_image_filename="ico_rgb_image.png"):
        """ Draw the ico's 3D tangent plan.

        :param radius: the radius of icosahedron.
        :type radius: float
        :param radius: the tangent image padding size
        :type radius: float
        :param obj_file_path: the obj file output path.
        :type obj_file_path: str
        """
        # 1) the panes center 3D center point
        if obj_file_texture_enable:
            mtl_file_path_path = Path(obj_file_path)
            mtl_file_path_path = mtl_file_path_path.with_suffix('.mtl')
            obj_file_head = "mtllib {}\n".format(str(mtl_file_path_path))
            obj_file_head += "usemtl material_0\n"
            obj_file_texture_coor_str = ""
        else:
            obj_file_head = ""

        obj_file_vertex_str = ""
        # the texture coordinate
        obj_file_face_str = ""
        obj_file_line_str = ""

        for tangent_image_idx in range(0, 20, 1):
            tangent_image_param = \
                pro_ico.get_icosahedron_parameters(tangent_image_idx, padding_size)

            triangle_points_sph = tangent_image_param["triangle_points_sph"]
            triangle_points_sph = np.array(triangle_points_sph, dtype=np.float32)

            triangle_points_sph_sort_idx = np.argsort(triangle_points_sph[:, 0])
            if (triangle_points_sph_sort_idx != np.array((0, 1, 2))).any():
                triangle_points_sph_sort_idx = np.tile(triangle_points_sph_sort_idx, 2).reshape((2, 3))
                triangle_points_sph = np.take_along_axis(triangle_points_sph, triangle_points_sph_sort_idx.T, axis=0)

            triangle_points_3d = sc.sph2car(triangle_points_sph[:, 0], triangle_points_sph[:, 1], radius=radius)
            # find the head vertex
            triangle_points_3d_head = None
            triangle_points_3d_edge = np.zeros((2, 3), dtype=np.float32)
            if triangle_points_sph[0, 1] == triangle_points_sph[1, 1]:
                triangle_points_3d_head = triangle_points_3d[:, 2]
                triangle_points_3d_edge[0, :] = triangle_points_3d[:, 0]
                triangle_points_3d_edge[1, :] = triangle_points_3d[:, 1]
            elif triangle_points_sph[0, 1] == triangle_points_sph[2, 1]:
                triangle_points_3d_head = triangle_points_3d[:, 1]
                triangle_points_3d_edge[0, :] = triangle_points_3d[:, 0]
                triangle_points_3d_edge[1, :] = triangle_points_3d[:, 2]
            else:
                triangle_points_3d_head = triangle_points_3d[:, 0]
                triangle_points_3d_edge[0, :] = triangle_points_3d[:, 1]
                triangle_points_3d_edge[1, :] = triangle_points_3d[:, 2]

            # 2) Output the vertices, texture information
            # the tangent plane rectangle tangent images
            if subimage_shift_ratio != 0.0:
                # offset the tangent images
                triangle_normal_norm_sph = tangent_image_param["tangent_point"]
                triangle_normal_norm_3d = sc.sph2car(triangle_normal_norm_sph[0], triangle_normal_norm_sph[1], radius=radius)
                subimage_shift = triangle_normal_norm_3d * subimage_shift_ratio
                triangle_points_3d_edge += subimage_shift
                triangle_points_3d_head += subimage_shift

            # 2-0) 3D points
            rectangle_points_3d = polygon.triangle_bounding_rectangle_3D(triangle_points_3d_head, triangle_points_3d_edge)
            # write vertex
            for vec_idx in range(4):
                vec_val = rectangle_points_3d[vec_idx]
                obj_file_vertex_str += "v {} {} {}\n".format(vec_val[0], vec_val[1], vec_val[2])

            # 2-1) add texture
            # write texture name and image file
            if obj_file_texture_enable:
                # texture file information
                # obj_file_face_str += "usemtl material_0\n"

                # texture offset
                texture_x_offset_start = 1.0 / 5.0 * (tangent_image_idx % 5)
                texture_x_offset_end = 1.0 / 5.0 + texture_x_offset_start
                texture_y_offset_start = 1.0 - 1.0 / 4.0 * int(tangent_image_idx / 5)
                texture_y_offset_end = texture_y_offset_start - 1.0 / 4.0
                obj_file_texture_coor_str += "vt {} {}\n".format(texture_x_offset_start, texture_y_offset_start)
                obj_file_texture_coor_str += "vt {} {}\n".format(texture_x_offset_start, texture_y_offset_end)
                obj_file_texture_coor_str += "vt {} {}\n".format(texture_x_offset_end, texture_y_offset_end)
                obj_file_texture_coor_str += "vt {} {}\n".format(texture_x_offset_end, texture_y_offset_start)

            # 2-2) write face
            # the head vertex alway at the top, make a face outside
            head_vertex_up = triangle_points_3d_head[1] < triangle_points_3d_edge[0, 1]
            # if head_vertex_up:
            #     triangle_points_3d_edge[[0, 1], :] = triangle_points_3d_edge[[1, 0], :]
            idx_offset = tangent_image_idx * 4
            if obj_file_texture_enable:
                if head_vertex_up:
                    print("{}:up".format(tangent_image_idx))
                    obj_file_face_str += "f {}/{} {}/{} {}/{}\n".format(idx_offset + 1, 1 + idx_offset,
                                                                        idx_offset + 2, 2 + idx_offset,
                                                                        idx_offset + 3, 3 + idx_offset)
                    obj_file_face_str += "f {}/{} {}/{} {}/{}\n".format(idx_offset + 1, 1 + idx_offset,
                                                                        idx_offset + 3, 3 + idx_offset,
                                                                        idx_offset + 4, 4 + idx_offset)
                else:
                    print("{}:down".format(tangent_image_idx))
                    obj_file_face_str += "f {}/{} {}/{} {}/{}\n".format(idx_offset + 3, 1 + idx_offset,
                                                                        idx_offset + 4, 2 + idx_offset,
                                                                        idx_offset + 1, 3 + idx_offset)
                    obj_file_face_str += "f {}/{} {}/{} {}/{}\n".format(idx_offset + 3, 1 + idx_offset,
                                                                        idx_offset + 1, 3 + idx_offset,
                                                                        idx_offset + 2, 4 + idx_offset)
            else:
                obj_file_face_str += "f {} {} {}\n".format(idx_offset + 1, idx_offset + 2, idx_offset + 3)
                obj_file_face_str += "f {} {} {}\n".format(idx_offset + 1, idx_offset + 3, idx_offset + 4)

            obj_file_line_str += "f {} {} {} {}\n".format(idx_offset + 1, idx_offset + 2, idx_offset + 3, idx_offset + 4)

        # 3) write obj file
        # 3-0) tangent images obj file
        # write tangent image obj file
        log.info("output the mesh to {}".format(obj_file_path))
        with open(obj_file_path, "w") as obj_file:
            obj_file.write(obj_file_head)
            obj_file.write(obj_file_vertex_str)
            if obj_file_texture_enable:
                obj_file.write(obj_file_texture_coor_str)

            obj_file.write(obj_file_face_str)

        # write tangent image mtl files
        if obj_file_texture_enable:
            # output mtl
            mtl_file_str = ""
            mtl_file_str += "newmtl material_0\n"
            mtl_file_str += "Ka 0.200000 0.200000 0.200000\n"
            mtl_file_str += "Kd 1.000000 1.000000 1.000000\n"
            mtl_file_str += "Ks 1.000000 1.000000 1.000000\n"
            mtl_file_str += "map_Kd {}\n\n".format(texture_image_filename)
            with open(str(mtl_file_path_path), "w") as mtl_file:
                mtl_file.write(mtl_file_str)

        # 3-1) write framewire obj file
        # write line obj
        obj_framewire_obj_file = obj_file_path + "_frameware.obj"
        log.info("output the framewire to {}".format(obj_file_path))
        with open(obj_framewire_obj_file, "w") as obj_file:
            obj_file.write(obj_file_vertex_str)
            obj_file.write(obj_file_line_str)


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

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems and the Max Planck Institute for Biological
# Cybernetics. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de

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

import time
try:
    import cPickle as pickle
except ImportError:
    import pickle

import sys
import os
import os.path as osp

import numpy as np
import torch

from tqdm import tqdm

from collections import defaultdict

import cv2
import PIL.Image as pil_img

from optimizers import optim_factory

import fitting
from vposer.model_loader import load_vposer

def fit_single_frame(img_list,
                     keypoints_list,
                     body_model,
                     camera_list,
                     joint_weights,
                     body_pose_prior,
                     jaw_prior,
                     left_hand_prior,
                     right_hand_prior,
                     shape_prior,
                     expr_prior,
                     angle_prior,
                     output_folder='',
                     result_fn='out.pkl',
                     mesh_fn='out.obj',
                     loss_type='smplify',
                     use_cuda=True,
                     init_joints_idxs=(9, 12, 2, 5),
                     use_face=True,
                     use_hands=True,
                     data_weights=None,
                     body_pose_prior_weights=None,
                     hand_pose_prior_weights=None,
                     jaw_pose_prior_weights=None,
                     shape_weights=None,
                     expr_weights=None,
                     hand_joints_weights=None,
                     face_joints_weights=None,
                     depth_loss_weight=1e2,
                     interpenetration=True,
                     coll_loss_weights=None,
                     df_cone_height=0.5,
                     penalize_outside=True,
                     max_collisions=8,
                     point2plane=False,
                     part_segm_fn='',
                     focal_length=5000.,
                     side_view_thsh=25.,
                     rho=100,
                     vposer_latent_dim=32,
                     vposer_ckpt='',
                     use_joints_conf=False,
                     interactive=True,
                     visualize=False,
                     save_meshes=True,
                     degrees=None,
                     batch_size=1,
                     dtype=torch.float32,
                     ign_part_pairs=None,
                     left_shoulder_idx=2,
                     right_shoulder_idx=5,
                     **kwargs):
    assert batch_size == 1, 'PyTorch L-BFGS only supports batch_size == 1'
    assert len(img_list) == len(keypoints_list)
    assert len(img_list) == len(camera_list)

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

    if degrees is None:
        degrees = [0, 90, 180, 270]

    if data_weights is None:
        data_weights = [1, ] * 5

    if body_pose_prior_weights is None:
        body_pose_prior_weights = [4.04 * 1e2, 4.04 * 1e2, 57.4, 4.78]

    msg = (
        'Number of Body pose prior weights {}'.format(
            len(body_pose_prior_weights)) +
        ' does not match the number of data term weights {}'.format(
            len(data_weights)))
    assert (len(data_weights) ==
            len(body_pose_prior_weights)), msg

    if use_hands:
        if hand_pose_prior_weights is None:
            hand_pose_prior_weights = [1e2, 5 * 1e1, 1e1, .5 * 1e1]
        msg = ('Number of Body pose prior weights does not match the' +
               ' number of hand pose prior weights')
        assert (len(hand_pose_prior_weights) ==
                len(body_pose_prior_weights)), msg
        if hand_joints_weights is None:
            hand_joints_weights = [0.0, 0.0, 0.0, 1.0]
            msg = ('Number of Body pose prior weights does not match the' +
                   ' number of hand joint distance weights')
            assert (len(hand_joints_weights) ==
                    len(body_pose_prior_weights)), msg

    if shape_weights is None:
        shape_weights = [1e2, 5 * 1e1, 1e1, .5 * 1e1]
    msg = ('Number of Body pose prior weights = {} does not match the' +
           ' number of Shape prior weights = {}')
    assert (len(shape_weights) ==
            len(body_pose_prior_weights)), msg.format(
                len(shape_weights),
                len(body_pose_prior_weights))

    if use_face:
        if jaw_pose_prior_weights is None:
            jaw_pose_prior_weights = [[x] * 3 for x in shape_weights]
        else:
            jaw_pose_prior_weights = map(lambda x: map(float, x.split(',')),
                                         jaw_pose_prior_weights)
            jaw_pose_prior_weights = [list(w) for w in jaw_pose_prior_weights]
        msg = ('Number of Body pose prior weights does not match the' +
               ' number of jaw pose prior weights')
        assert (len(jaw_pose_prior_weights) ==
                len(body_pose_prior_weights)), msg

        if expr_weights is None:
            expr_weights = [1e2, 5 * 1e1, 1e1, .5 * 1e1]
        msg = ('Number of Body pose prior weights = {} does not match the' +
               ' number of Expression prior weights = {}')
        assert (len(expr_weights) ==
                len(body_pose_prior_weights)), msg.format(
                    len(body_pose_prior_weights),
                    len(expr_weights))

        if face_joints_weights is None:
            face_joints_weights = [0.0, 0.0, 0.0, 1.0]
        msg = ('Number of Body pose prior weights does not match the' +
               ' number of face joint distance weights')
        assert (len(face_joints_weights) ==
                len(body_pose_prior_weights)), msg

    if coll_loss_weights is None:
        coll_loss_weights = [0.0] * len(body_pose_prior_weights)
    msg = ('Number of Body pose prior weights does not match the' +
           ' number of collision loss weights')
    assert (len(coll_loss_weights) ==
            len(body_pose_prior_weights)), msg

    use_vposer = kwargs.get('use_vposer', True)
    vposer, pose_embedding = [None, ] * 2
    if use_vposer:
        pose_embedding = torch.zeros([batch_size, 32],
                                     dtype=dtype, device=device,
                                     requires_grad=True)

        vposer_ckpt = osp.expandvars(vposer_ckpt)
        vposer, _ = load_vposer(vposer_ckpt, vp_model='snapshot')
        vposer = vposer.to(device=device)
        vposer.eval()

    if use_vposer:
        body_mean_pose = torch.zeros([batch_size, vposer_latent_dim],
                                     dtype=dtype)
    else:
        body_mean_pose = body_pose_prior.get_mean().detach().cpu()

    view_num = len(camera_list)
    loss_list = list()
    gt_joints_list = list()
    joints_conf_list = list()

    assert(view_num > 0)
    for view_id in range(view_num):
        keypoint_data = torch.tensor(keypoints_list[view_id], dtype=dtype)
        gt_joints = keypoint_data[:, :, :2]
        if use_joints_conf:
            joints_conf = keypoint_data[:, :, 2].reshape(1, -1)

        # Transfer the data to the correct device
        gt_joints = gt_joints.to(device=device, dtype=dtype)
        gt_joints_list.append(gt_joints)
        if use_joints_conf:
            joints_conf = joints_conf.to(device=device, dtype=dtype)
            joints_conf_list.append(joints_conf)

        # Create the search tree
        search_tree = None
        pen_distance = None
        filter_faces = None
        if interpenetration:
            raise NotImplementedError('The interpenetration constraint was removed!')

        fct = view_num

        # Weights used for the pose prior and the shape prior
        opt_weights_dict = {'data_weight': data_weights,
                            'body_pose_weight': body_pose_prior_weights,
                            'shape_weight': shape_weights}
        # adjust energy weight for multi-view setup
        for i in range(len(opt_weights_dict['body_pose_weight'])):
            opt_weights_dict['body_pose_weight'][i] *= (view_num / fct)
        for i in range(len(opt_weights_dict['shape_weight'])):
            opt_weights_dict['shape_weight'][i] *= (view_num / fct)

        if use_face:
            opt_weights_dict['face_weight'] = face_joints_weights
            opt_weights_dict['expr_prior_weight'] = expr_weights
            opt_weights_dict['jaw_prior_weight'] = jaw_pose_prior_weights
            for i in range(len(opt_weights_dict['expr_prior_weight'])):
                opt_weights_dict['expr_prior_weight'][i] *= (view_num / fct)
            for i in range(len(opt_weights_dict['jaw_prior_weight'])):
                opt_weights_dict['jaw_prior_weight'][i] *= (view_num / fct)

        if use_hands:
            opt_weights_dict['hand_weight'] = hand_joints_weights
            opt_weights_dict['hand_prior_weight'] = hand_pose_prior_weights
            for i in range(len(opt_weights_dict['hand_prior_weight'])):
                opt_weights_dict['hand_prior_weight'][i] *= (view_num / fct)

        if interpenetration:
            opt_weights_dict['coll_loss_weight'] = coll_loss_weights
            for i in range(len(opt_weights_dict['coll_loss_weight'])):
                opt_weights_dict['coll_loss_weight'][i] *= (view_num / fct)

        keys = opt_weights_dict.keys()
        opt_weights = [dict(zip(keys, vals)) for vals in
                       zip(*(opt_weights_dict[k] for k in keys
                             if opt_weights_dict[k] is not None))]
        for weight_list in opt_weights:
            for key in weight_list:
                weight_list[key] = torch.tensor(weight_list[key],
                                                device=device,
                                                dtype=dtype)

        loss = fitting.create_loss(loss_type=loss_type,
                                   joint_weights=joint_weights,
                                   rho=rho,
                                   use_joints_conf=use_joints_conf,
                                   use_face=use_face, use_hands=use_hands,
                                   vposer=vposer,
                                   pose_embedding=pose_embedding,
                                   body_pose_prior=body_pose_prior,
                                   shape_prior=shape_prior,
                                   angle_prior=angle_prior,
                                   expr_prior=expr_prior,
                                   left_hand_prior=left_hand_prior,
                                   right_hand_prior=right_hand_prior,
                                   jaw_prior=jaw_prior,
                                   interpenetration=interpenetration,
                                   pen_distance=pen_distance,
                                   search_tree=search_tree,
                                   tri_filtering_module=filter_faces,
                                   # I scale the mesh model to [-0.5, 0.5] during rendering;
                                   # so I need to perform the same scaling
                                   # to make the body shapee plausible
                                   dtype=dtype,
                                   **kwargs)
        loss = loss.to(device=device)
        loss_list.append(loss)

    body_scale = torch.tensor([1.0 / 1.7], dtype=dtype, device=device,
                              requires_grad=True)
    global_body_translation = torch.tensor([0, 0, 0], dtype=dtype, device=device,
                                           requires_grad=True)

    with fitting.FittingMonitor(
            batch_size=batch_size, visualize=visualize, **kwargs) as monitor:

        img = torch.tensor(img_list[0], dtype=dtype)

        H, W, _ = img.shape

        data_weight = 1000 / H

        # Reset the parameters to estimate the initial translation of the
        # body model
        body_model.reset_params(body_pose=body_mean_pose)

        # If the distance between the 2D shoulders is smaller than a
        # predefined threshold then try 2 fits, the initial one and a 180
        # degree rotation
        # shoulder_dist = torch.dist(gt_joints[:, left_shoulder_idx],
        #                            gt_joints[:, right_shoulder_idx])
        # try_both_orient = shoulder_dist.item() < side_view_thsh
        try_both_orient = False

        # If the 2D detections/positions of the shoulder joints are too
        # close the rotate the body by 180 degrees and also fit to that
        # orientation
        if try_both_orient:
            body_orient = body_model.global_orient.detach().cpu().numpy()
            flipped_orient = cv2.Rodrigues(body_orient)[0].dot(
                cv2.Rodrigues(np.array([0., np.pi, 0]))[0])
            flipped_orient = cv2.Rodrigues(flipped_orient)[0].ravel()

            flipped_orient = torch.tensor(flipped_orient,
                                          dtype=dtype,
                                          device=device).unsqueeze(dim=0)
            orientations = [body_orient, flipped_orient]
        else:
            orientations = [body_model.global_orient.detach().cpu().numpy()]

        # store here the final error for both orientations,
        # and pick the orientation resulting in the lowest error
        results = []

        # Step 2: Optimize the full model
        final_loss_val = 0
        for or_idx, orient in enumerate(tqdm(orientations, desc='Orientation')):
            opt_start = time.time()

            new_params = defaultdict(global_orient=orient,
                                     body_pose=body_mean_pose)
            body_model.reset_params(**new_params)
            if use_vposer:
                with torch.no_grad():
                    pose_embedding.fill_(0)

            for opt_idx, curr_weights in enumerate(tqdm(opt_weights, desc='Stage')):

                body_params = list(body_model.parameters())

                final_params = list(
                    filter(lambda x: x.requires_grad, body_params))

                if use_vposer:
                    final_params.append(pose_embedding)
                final_params.append(global_body_translation)
                final_params.append(body_scale)
                body_optimizer, body_create_graph = optim_factory.create_optimizer(
                    final_params,
                    **kwargs)
                body_optimizer.zero_grad()

                curr_weights['data_weight'] = data_weight
                curr_weights['bending_prior_weight'] = (
                    3.17 * curr_weights['body_pose_weight'])
                if use_hands:
                    joint_weights[:, 25:76] = curr_weights['hand_weight']
                if use_face:
                    joint_weights[:, 76:] = curr_weights['face_weight']

                for i in range(len(loss_list)):
                    loss_list[i].reset_loss_weights(curr_weights)

                closure = monitor.create_fitting_closure_multiview(
                    body_optimizer, body_model,
                    camera_list=camera_list, global_body_translation=global_body_translation,
                    body_model_scale=body_scale,
                    gt_joints_list=gt_joints_list,
                    joints_conf_list=joints_conf_list,
                    joint_weights=joint_weights,
                    loss_list=loss_list, create_graph=body_create_graph,
                    use_vposer=use_vposer, vposer=vposer,
                    pose_embedding=pose_embedding,
                    return_verts=True, return_full_pose=True)

                if interactive:
                    if use_cuda and torch.cuda.is_available():
                        torch.cuda.synchronize()
                    stage_start = time.time()
                final_loss_val = monitor.run_fitting(
                    body_optimizer,
                    closure, final_params,
                    body_model,
                    pose_embedding=pose_embedding, vposer=vposer,
                    use_vposer=use_vposer)

                if interactive:
                    if use_cuda and torch.cuda.is_available():
                        torch.cuda.synchronize()
                    elapsed = time.time() - stage_start
                    if interactive:
                        tqdm.write('Stage {:03d} done after {:.4f} seconds'.format(
                            opt_idx, elapsed))

            if interactive:
                if use_cuda and torch.cuda.is_available():
                    torch.cuda.synchronize()
                elapsed = time.time() - opt_start
                tqdm.write(
                    'Body fitting Orientation {} done after {:.4f} seconds'.format(
                        or_idx, elapsed))
                tqdm.write('Body final loss val = {:.5f}'.format(
                    final_loss_val))

            # Get the result of the fitting process
            # Store in it the errors list in order to compare multiple
            # orientations, if they exist
            # result = {'camera_' + str(key): val.detach().cpu().numpy()
            #           for key, val in camera.named_parameters()}
            result = {}
            result.update({key: val.detach().cpu().numpy()
                           for key, val in body_model.named_parameters()})
            result.update({'global_body_translation':
                                global_body_translation.detach().cpu().numpy()})
            result.update({'body_scale':
                               body_scale.detach().cpu().numpy()})
            if use_vposer:
                body_pose = vposer.decode(
                    pose_embedding,
                    output_type='aa').view(1, -1) if use_vposer else None

                model_type = kwargs.get('model_type', 'smpl')
                append_wrists = model_type == 'smpl' and use_vposer
                if append_wrists:
                    wrist_pose = torch.zeros([body_pose.shape[0], 6],
                                             dtype=body_pose.dtype,
                                             device=body_pose.device)
                    body_pose = torch.cat([body_pose, wrist_pose], dim=1)

                result['body_pose'] = body_pose.detach().cpu().numpy()
                result['body_pose_embedding'] = pose_embedding.detach().cpu().numpy()

            if result['body_pose'].shape[-1] == 69:
                body_pose = result['body_pose']
                body_pose = np.reshape(body_pose, (1, 69))
                body_pose = np.concatenate([result['global_orient'], body_pose], axis=1)
                result.update({'body_pose': body_pose})

            assert result['body_pose'].shape[-1] == 72
            results.append({'loss': final_loss_val,
                            'result': result})
            print('body_scale = %f' % body_scale.detach().cpu().numpy().squeeze())

        with open(result_fn, 'wb') as result_file:
            if len(results) > 1:
                min_idx = (0 if results[0]['loss'] < results[1]['loss']
                           else 1)
            else:
                min_idx = 0
            pickle.dump(results[min_idx]['result'], result_file, protocol=2)

    if save_meshes or visualize:
        model_output = body_model(return_verts=True, body_pose=torch.from_numpy(result['body_pose'][:, 3:]).cuda())
        vertices = model_output.vertices.detach().cpu().numpy().squeeze()

        # test projection
        global_trans = global_body_translation.detach().cpu().numpy().squeeze()
        body_scale = body_scale.detach().cpu().numpy().squeeze()

        # project smpl vertices onto images for debugging
        out_img_fd = osp.join(output_folder, 'vis')
        os.makedirs(out_img_fd, exist_ok=True)
        for i, (camera, img) in enumerate(zip(camera_list, img_list)):
            cam_fx = camera.focal_length_x.detach().cpu().numpy().squeeze()
            cam_fy = camera.focal_length_y.detach().cpu().numpy().squeeze()
            cam_c = camera.center.detach().cpu().numpy().squeeze()
            cam_trans = camera.translation.detach().cpu().numpy().squeeze()
            cam_rotation = camera.rotation.detach().cpu().numpy().squeeze()

            vertices_proj = vertices * body_scale + global_trans
            vertices_proj = np.dot(vertices_proj, cam_rotation.transpose())
            vertices_proj += np.expand_dims(cam_trans, axis=0)
            vertices_proj[:, 0] = vertices_proj[:, 0] * cam_fx / vertices_proj[:, 2] + cam_c[0]
            vertices_proj[:, 1] = vertices_proj[:, 1] * cam_fy / vertices_proj[:, 2] + cam_c[1]
            img_proj = np.copy(img)
            for v in vertices_proj:
                v = np.int32(np.round(v))
                v[0] = np.clip(v[0], 0, img_proj.shape[1]-1)
                v[1] = np.clip(v[1], 0, img_proj.shape[0]-1)
                img_proj[v[1], v[0], :] = np.asarray([0, 0, 1], dtype=np.float32)
            img_proj = np.uint8(img_proj*255)
            cv2.imwrite(osp.join(out_img_fd, '%04d.png' % i), img_proj)

        import trimesh
        out_mesh = trimesh.Trimesh(vertices * body_scale + global_trans, body_model.faces)
        out_mesh.export(mesh_fn)