import dataProcessing as dp
import plotFuncs as pf

import numpy as np
from matplotlib.ticker import MultipleLocator, FormatStrFormatter
from mpl_toolkits.axes_grid1 import host_subplot
import mpl_toolkits.axisartist as AA

from matplotlib.path import Path
from mpl_toolkits.mplot3d import Axes3D
import matplotlib as mpl

from compAnalysisBase import CompAnalysisBase

class SIMCompAnalysis(CompAnalysisBase):

    def __init__(self, leavesQueue, criteria, sifs):
        self.queue = leavesQueue
        self.sifs = sifs
        self.crit = criteria

    def printQueueItems(self, items):
        self.queue.printTitle()
        for i in sorted(items):
            self.queue.printQueueItem(i)

    def getItemNodeDict(self, items, queue):
        qdict = queue.getQueueDict()
        return dict([(i, qdict[i]) for i in items])

    def calcAlphaVal(self, sif, item):
        vals = len(self.dataDicts[0][0][sif][item])
        if vals > 1000:
            return 0.1
        else:
            return 1

class BoxCompPlot(SIMCompAnalysis):

    def createCompBoxPlot(self, items, errType, fig):
        self.items = items
        self.errType = errType
        self.createDataDictAndEstBoxPlot()
        self.createDataStrBoxPlot()
        self.createFigure(fig)

    def createDataStrBoxPlot(self):
        dd = self.getItemNodeDict(self.items, self.queue)
        optKey = self.getLeavesOptKey()
        data = [dd, optKey, 'Number in Queue', '']
        self.dataStr = [data]

    def getLeavesOptKey(self):
        return sorted(self.est.items(), key=lambda x: abs(x[1]))[0][0]

    def createDataDictAndEstBoxPlot(self):
        dataDict = {s: {} for s in self.sifs}
        est = {i: {} for i in self.items}
        dd = self.getItemNodeDict(self.items, self.queue)
        for i in self.items:
            node = dd[i]
            errs, est[i] = self.getNodeErrsEst(node)
            for s in self.sifs:
                dataDict[s][i] = errs[s]
        self.est = {i: est[i][self.crit[1]] for i in self.items}
        self.dataDicts = [dataDict]

    def getNodeErrsEst(self, node):
        adn = dp.AnalysisNodeData(node, self.sifs)
        adn.performOperations()
        est = adn.getEstimates()[self.crit[0]]
        errs = adn.getErrors()[self.errType]
        return errs, est

class HistCompPlot(SIMCompAnalysis):

    def createCompHistPlot(self, items, errType, xlim, fig):
        self.fig = fig
        self.items = items
        self.errType = errType
        self.xlim = xlim
        self.createDataStr()
        self.createDataDict()
        self.createFigure()

    def createDataStr(self):
        dd = self.getItemNodeDict(self.items.keys(), self.queue)
        xlabel = 'errors "{0}"'.format(self.errType)
        data = [dd, None, xlabel, 'hist']
        self.dataStr = [data]

    def createDataDict(self):
        data = {s: {} for s in self.sifs}
        for i in self.items.keys():
            node = self.dataStr[0][0][i]
            errs = self.getNodeErrors(node)
            for s in self.sifs:
                data[s][i] = errs[s]
        self.dataDicts = [data]

    def getNodeErrors(self, node):
        adn = dp.AnalysisNodeData(node, self.sifs)
        adn.performOperations()
        errs = adn.getErrors()[self.errType]
        return errs

    def setAxesXlim(self):
        for ax in self.axes:
            ax.set_xlim(self.xlim)

    def setAxesYlim(self):
        ymin, ymax = 10e16, 10e-16
        for ax in self.axes:
            y1, y2 = ax.get_ylim()
            ymin = y1 if y1 < ymin else ymin
            ymax = y2 if y2 > ymax else ymax
        for ax in self.axes:
            ax.set_ylim((ymin, ymax))

    def setLegend(self, handles):
        text = 'Node: '
        labels = [text + str(i) for i in sorted(handles.keys())]
        handles = [handles[i] for i in sorted(handles.keys())]
        self.axes[0].legend(handles, labels, bbox_to_anchor=(1.02, 1),
                            loc=2, borderaxespad=0)

    def createFigure(self):
        self.axes = []
        self.createFigureAxes()
        handles = {}
        for k in range(len(self.axes)):
            s = self.sifs[k]
            for i in self.items.keys():
                n, b, p = self.axes[k].hist(
                    self.dataDicts[0][s][i],
                    self.items[i],
                    normed=True,
                    alpha=0.5)
                handles[i] = p[0]
        self.setAxesXlim()
        self.setAxesYlim()
        self.setLegend(handles)
        self.setXlabels()
        self.printQueueItems(self.items.keys())

class CorrCompPlot(SIMCompAnalysis):

    def createCompCorrPlot(self, items, quantityType, ylim, fig):
        self.fig = fig
        self.items = items
        self.qt = quantityType
        self.ylim = ylim
        self.createDataStr()
        self.createDataDict()
        self.createFigure()

    def createDataStr(self):
        dd = self.getItemNodeDict(self.items, self.queue)
        data = [dd, None, 'analytical values', 'analysis vs analytical']
        self.dataStr = [data]

    def createDataDict(self):
        dataX = {s: {} for s in self.sifs}
        dataY = {s: {} for s in self.sifs}
        for i in self.items:
            node = self.dataStr[0][0][i]
            anSol, res = self.getNodeParams(node)
            for s in self.sifs:
                dataX[s][i] = anSol[s]
                dataY[s][i] = res[s]
        self.dataDicts = [[dataX, dataY]]

    def getNodeParams(self, node):
        adn = dp.AnalysisNodeData(node, self.sifs)
        adn.performOperations()
        anSol = adn.getAnSol()
        res = adn.getDataByType(self.qt)
        return anSol, res

    def getReferenceXYVals(self):
        minV = {s: 10e16 for s in self.sifs}
        maxV = {s: -10e16 for s in self.sifs}
        for s in self.sifs:
            for i in self.items:
                mn = min(self.dataDicts[0][0][s][i])
                mx = max(self.dataDicts[0][0][s][i])
                minV[s] = mn if mn < minV[s] else minV[s]
                maxV[s] = mx if mx > maxV[s] else maxV[s]
        if self.qt == 'results':
            refX = {s: [minV[s], maxV[s]] for s in self.sifs}
            return refX, refX
        elif self.qt in ['difference', 'normedDiff']:
            refX = {s: [max(0, minV[s]), maxV[s]] for s in self.sifs}
            refY = {s: [0, 0] for s in self.sifs}
            return refX, refY
        else:
            raise NotImplementedError

    def getXYVals(self, sif, item):
        if self.qt == 'results':
            X = self.dataDicts[0][0][sif][item]
            Y = self.dataDicts[0][1][sif][item]
        elif self.qt in ['difference', 'normedDiff']:
            X = np.abs(self.dataDicts[0][0][sif][item])
            Y = self.dataDicts[0][1][sif][item]
        else:
            raise NotImplementedError
        return X, Y

    def createPlot(self):
        self.handles = {}
        refX, refY = self.getReferenceXYVals()
        for k in range(len(self.axes)):
            s = self.sifs[k]
            for i in self.items:
                alpha = self.calcAlphaVal(s, i)
                X, Y = self.getXYVals(s, i)
                p, = self.axes[k].plot(X, Y, '.', alpha=alpha)
                self.handles[i] = p
            r, = self.axes[k].plot(refX[s], refY[s], 'k', lw=1.5)
        self.handles['reference'] = r

    def setXLim(self):
        refX, refY = self.getReferenceXYVals()
        for k in range(len(self.axes)):
            s = self.sifs[k]
            self.axes[k].set_xlim(refX[s])

    def setLegend(self):
        text = 'Node: '
        labels = [text + str(i) for i in self.items]
        handles = [self.handles[i] for i in self.items]
        if 'reference' in self.handles.keys():
            handles.append(self.handles['reference'])
            labels.append('ref line')
        self.axes[0].legend(handles, labels, bbox_to_anchor=(1.02, 1),
                            loc=2, borderaxespad=0)

    def setYLim(self):
        if isinstance(self.ylim, (list, tuple)):
            for ax in self.axes:
                ax.set_ylim(self.ylim)

    def createFigure(self):
        self.axes = []
        self.createFigureAxes()
        self.createPlot()
        self.setXLim()
        self.setLegend()
        self.printQueueItems(self.items)
        self.setYLim()

class RangeCompPlot(SIMCompAnalysis):

    def createCompRangePlot(self, items, opts, fig):
        self.fig = fig
        self.items = items
        self.opts = opts
        self.createDataStr()
        self.createDataDict()
        self.createFigure()

    def createDataStr(self):
        self.dataStr = []
        qdict = self.queue.getQueueDict()
        for k in sorted(self.items.keys()):
            optSim = self.getOptSim(qdict[k])
            data = [{k: qdict[k]}, optSim, 'angles',
                    self.getSubplotTitle(qdict[k])]
            self.dataStr.append(data)

    def getOptSim(self, node):
        if self.opts['optSim']:
            sims = node.getSuccessfulMembers()
            optSim = pf.getSimIdsWithLowestErrorPerDH(
                sims, self.crit[0], self.crit[1]).values()[0][0]
            return optSim
        else:
            return None

    def createDataDict(self):
        self.dataDicts = []
        for item in self.dataStr:
            node = item[0].values()[0]
            self.dataDicts.append(self.getNodeParams(node))

    def getNodeParams(self, node):
        adn = dp.AnalysisNodeData(node, self.sifs)
        adn.performOperations()
        angles = adn.getAngles()
        results = adn.getResults()
        ansol = adn.getAnSol()
        errors = adn.getErrors()[self.opts['errors']]
        return angles, results, ansol, errors

    def createSlices(self):
        self.slices = []
        i = 0
        for k in sorted(self.items.keys()):
            numInt = self.items[k]
            angles = self.dataDicts[i][0]
            sl = self.createSliceIndices(angles, numInt)
            self.slices.append(sl)
            i += 1

    def createSliceIndices(self, vals, numInts):
        intLen = (max(vals) - min(vals)) / float(numInts)
        indices = [[] for i in range(numInts)]
        for x in vals:
            i = int(x / intLen)
            if i < numInts - 1:
                indices[i].append(x)
            else:
                indices[-1].append(x)
        if [] in indices:
            raise ValueError('Try reducing the number of intervals.')
        sliceInd = [[] for i in range(numInts)]
        for i in range(numInts):
            minVal = indices[i][0]
            maxVal = indices[i][-1]
            ind0 = np.where(vals == minVal)[0][0]
            ind1 = np.where(vals == maxVal)[-1][-1] + 1
            sliceInd[i].append(ind0)
            sliceInd[i].append(ind1)
        sliceInd[-1][1] += 1
        return sliceInd

    def createFigure(self):
        self.axes = []
        self.createFigureAxes()
        if self.opts['range']:
            self.createSlices()
            self.plotRangeArea()
        if self.opts['dataPoints']:
            self.createDataPointsPlot()
        if self.opts['analytical']:
            self.createAnSolPlot()
        if self.opts['optSim']:
            self.createOptSimPlot()
        self.setXLim()
        self.createLegend()
        self.setSubplotTitles()
        self.setYlimits()

    def createLegend(self):
        handles = []
        labels = []
        h, l = self.axes[0].get_legend_handles_labels()
        ind = len(self.dataStr) - 1
        self.axes[ind].legend(h, l, bbox_to_anchor=(1, 1.02), loc=2)

    def setXLim(self):
        for n in range(len(self.dataStr)):
            i = self.getItemKey(n)
            for sif in self.sifs:
                ax = self.getAxes(i, sif)
                angles = self.dataDicts[n][0]
                ax.set_xlim((min(angles), max(angles)))

    def createOptSimPlot(self):
        for n in range(len(self.dataDicts)):
            i = self.getItemKey(n)
            ad = dp.AnalysisData(self.dataStr[n][1])
            ad.calcAnSol()
            ad.calculateStats()
            angles = ad.getAngles()
            for sif in self.sifs:
                ax = self.getAxes(i, sif)
                res = ad.getResults()[sif]
                ax.plot(angles, res, 'lime', lw=1,
                        label='optSim')

    def createDataPointsPlot(self):
        for n in range(len(self.dataStr)):
            i = self.getItemKey(n)
            for sif in self.sifs:
                angles = self.dataDicts[n][0]
                ax = self.getAxes(i, sif)
                for dt in self.opts['data']:
                    dInd, color = self.getDataIndAndColor(dt)
                    data = self.dataDicts[n][dInd][sif]
                    alpha = self.calcAlphaValRP(n)
                    ax.plot(angles, data,
                            linestyle='-', marker='.',
                            color=color, alpha=alpha,
                            label=dt)

    def calcAlphaValRP(self, n):
        vals = len(self.dataDicts[n][0])
        if vals > 1000:
            return 0.05
        else:
            return 0.3

    def createAnSolPlot(self):
        for n in range(len(self.items.keys())):
            i = self.getItemKey(n)
            for sif in self.sifs:
                ax = self.getAxes(i, sif)
                angles = self.dataDicts[n][0]
                anSol = self.dataDicts[n][2][sif]
                ax.plot(angles, anSol, 'k', lw=2,
                        label='analytical')

    def getAxes(self, item, sif):
        itemInd = sorted(self.items.keys()).index(item)
        itemLen = len(self.items)
        ax = self.axes[itemLen * self.sifs.index(sif) + itemInd]
        return ax

    def getItemKey(self, n):
        return sorted(self.items.keys())[n]

    def plotRangeArea(self):
        for n in range(len(self.items)):
            i = self.getItemKey(n)
            for sif in self.sifs:
                axes = self.getAxes(i, sif)
                self.plotRangeAreaPerAxes(axes, n, sif)

    def getDataIndAndColor(self, dataType):
        dataInds = {'results': 1, 'errors': 3}
        colors = {'results': 'b', 'errors': 'r'}
        return dataInds[dataType], colors[dataType]

    def createVerts(self, slices, angles, values, func):
        x, y, verts = [], [], []
        valsl = [values[s[0] - 1 if s[0] > 0 else 0:s[1]] for s in slices]
        angsl = [angles[s[0] - 1 if s[0] > 0 else 0:s[1]] for s in slices]
        for a in angsl:
            x.append(a[0])
            x.append(a[-1])
        for v in valsl:
            y.append(func(v))
            y.append(func(v))
        verts = [[xi, yi] for xi, yi in zip(x, y)]
        return verts

    def createVerts2(self, slices, angles, values, func):
        x, y, verts = [], [], []
        valsl = [values[s[0]:s[1]] for s in slices]
        angsl = [angles[s[0]:s[1]] for s in slices]
        for an, va in zip(angsl, valsl):
            y.append(func(va))
            print va, y
            print np.where(va == y[-1])
            ind = np.where(va == y[-1])[0][0]
            x.append(an[ind])
        x.append(angles[-1])
        x.insert(0, angles[0])
        yavg = 0.5 * (y[0] + y[-1])
        y.append(yavg)
        y.insert(0, yavg)
        verts = [[xi, yi] for xi, yi in zip(x, y)]
        return verts

    def plotRangeAreaPerAxes(self, axes, itemInd, sif):
        vertMethods = {1: self.createVerts, 2: self.createVerts2}
        vertFunc = vertMethods[self.opts['rangeType']]
        slices = self.slices[itemInd]
        angles = self.dataDicts[itemInd][0]
        for dt in self.opts['data']:
            dInd, color = self.getDataIndAndColor(dt)
            values = self.dataDicts[itemInd][dInd][sif]
            verts1 = vertFunc(slices, angles, values, min)
            verts2 = vertFunc(slices, angles, values, max)[::-1]
            verts = verts1 + verts2 + [verts2[-1]]
            codes = self.createClosedPathCodes(verts)
            p = Path(verts, codes)
            patch = mpl.patches.PathPatch(
                p,
                facecolor=color,
                edgecolor='none',
                alpha=0.2,
                label=dt +
                ' range')
            axes.add_patch(patch)
            patch = mpl.patches.PathPatch(p, edgecolor=color,
                                          fill=False, lw=0.75, alpha=0.6)
            axes.add_patch(patch)

    def createClosedPathCodes(self, verts):
        codes = [Path.MOVETO]
        for i in range(len(verts) - 2):
            codes.append(Path.LINETO)
        codes.append(Path.CLOSEPOLY)
        return codes

class BoundsCompPlot(SIMCompAnalysis):

    def createBoundsPlot(self, items, targets, fig, tol=0.1, iterLim=100):
        self.items = items
        self.targets = targets
        self.fig = fig
        self.iterLim = iterLim
        self.tol = tol
        self.createDataStr()
        self.createDataDicts()
        self.printStats()
        self.createFigure()

    def createDataStr(self):
        self.dataStr = []
        qdict = self.queue.getQueueDict()
        for i in self.items:
            dd = [{i: qdict[i]}, None, 'angles',
                  self.getSubplotTitle(qdict[i])]
            self.dataStr.append(dd)

    def createDataDicts(self):
        self.dataDicts = []
        for n in range(len(self.items)):
            i = self.items[n]
            log = {s: {t: {'sigma': [], 'pip': []}
                       for t in self.targets.keys()}
                   for s in self.sifs}
            node = self.dataStr[n][0][i]
            adn = dp.AnalysisNodeData(node, self.sifs)
            adn.performOperations()
            sigmaUp = 2 * adn.getAnSolParams()['sigma']
            sigmaLow = 0
            for s in self.sifs:
                for t in self.targets.keys():
                    log[s][t] = self.findSigmaBound(
                        adn, sigmaUp, sigmaLow, s,
                        self.targets[t], log[s][t])
            self.dataDicts.append([adn, log])

    def printStats(self):
        for n in range(len(self.dataStr)):
            i = self.items[n]
            print self.dataStr[n][3]
            log = self.dataDicts[n][1]
            for s in self.sifs:
                sigmas, bounds, its = [], [], []
                for t in log[s].keys():
                    u = log[s][t]
                    sigmas.append(u['sigma'][-1])
                    bounds.append(u['pip'][-1])
                    its.append(len(u['sigma']))
                info = '{0}sigma=[{1:.4}, {2:.4}] | bounds=[{3:.4}%, {4:.4}%] | iterations=[{5}, {6}]'.format(
                    '  {0} '.format(s), sigmas[0], sigmas[1], bounds[0], bounds[1], its[0], its[1])
                print info

    def createFigure(self):
        self.axes = []
        self.createFigureAxes()
        self.createPlot()
        self.setXLimits()
        self.setYlimits()
        self.setSubplotTitles()

    def setXLimits(self):
        for n in range(len(self.dataStr)):
            i = self.items[n]
            adn = self.dataDicts[n][0]
            a = adn.getAngles()
            lims = (min(a), max(a))
            for s in self.sifs:
                ax = self.getAxes(i, s)
                ax.set_xlim(lims)

    def getAxes(self, item, sif):
        itemLen = len(self.items)
        itemInd = self.items.index(item)
        ax = self.axes[itemLen * self.sifs.index(sif) + itemInd]
        return ax

    def getAlphaVal(self, item):
        n = self.items.index(item)
        adn = self.dataDicts[n][0]
        if len(adn.getAngles()) > 1000:
            return 0.1
        else:
            return 1

    def createPlot(self):
        for n in range(len(self.dataStr)):
            i = self.items[n]
            adn = self.dataDicts[n][0]
            logs = self.dataDicts[n][1]
            alpha = self.getAlphaVal(i)
            for s in self.sifs:
                ax = self.getAxes(i, s)
                sigmaUpper = logs[s]['upper']['sigma'][-1]
                sigmaLower = logs[s]['lower']['sigma'][-1]
                ins, outs = self.getInOutPoints(adn,
                                                sigmaLower, sigmaUpper, s)
                ax.plot(ins[0], ins[1], 'b.',
                        label='inside bounds', alpha=alpha)
                ax.plot(outs[0], outs[1], 'r.',
                        label='outside bounds', alpha=alpha)
                angles = adn.getAngles()
                anSol = adn.getAnSol()[s]
                ax.plot(angles, anSol, 'k', lw=1.5,
                        label='analytical')
                lowerBound = adn.calcSIFsForSigmaAndSIF(
                    sigmaLower, s)
                upperBound = adn.calcSIFsForSigmaAndSIF(
                    sigmaUpper, s)
                ax.plot(angles, upperBound, 'lime', lw=1.5,
                        label='bounds')
                ax.plot(angles, lowerBound, 'lime', lw=1.5)

    def findSigmaBound(self, adn, sigmaUp, sigmaLow,
                       sif, target, log):
        sigma = 0.5 * (sigmaUp + sigmaLow)
        pip = self.getPercentPointsInPoly(adn, sigma, sif)
        log['pip'].append(pip)
        log['sigma'].append(sigma)
        if ((pip >= target - self.tol and pip <= target + self.tol) or
                (len(log['sigma']) == self.iterLim)):
            return log
        elif pip < target - self.tol:
            sigmaLow = sigma
            return self.findSigmaBound(adn, sigmaUp, sigmaLow,
                                       sif, target, log)
        elif pip > target + self.tol:
            sigmaUp = sigma
            return self.findSigmaBound(adn, sigmaUp, sigmaLow,
                                       sif, target, log)
        else:
            raise ValueError('unexpected condition reached')

    def getPercentPointsInPoly(self, adn, sigma, sif):
        allnum, numin, numout = self.countPointInOutOfContour(
            adn, sigma, sif)
        assert abs(numin + numout - allnum) < 10e-8
        return float(numin) / float(allnum) * 100

    def countPointInOutOfContour(self, adn, sigma, sif):
        tfl = self.getInOutOfContour(adn, sigma, sif)
        numin = np.sum(tfl)
        allnum = len(tfl)
        numout = allnum - numin
        return allnum, numin, numout

    def getInOutOfContour(self, adn, sigma, sif):
        angles = adn.getAngles()
        results = abs(adn.getResults()[sif])
        points = [[xi, yi] for xi, yi in zip(angles, results)]
        yVals = abs(np.array(adn.calcSIFsForSigmaAndSIF(sigma, sif)))
        return self.getInOutPointsArray(angles, yVals, points)

    def getInOutPointsArray(self, angles, yVals, points):
        path = Path(self.createVertsForPolyPath(angles, yVals))
        return path.contains_points(points, radius=0)

    def getInOutPoints(self, adn, sigmaLow, sigmaUp, sif):
        inoutLow = self.getInOutOfContour(adn, sigmaLow, sif)
        inoutUp = self.getInOutOfContour(adn, sigmaUp, sif)
        angles = adn.getAngles()
        res = adn.getResults()[sif]
        inAngles, inVals = [], []
        outAngles, outVals = [], []
        for i in range(len(inoutUp)):
            if inoutLow[i] or not inoutUp[i]:
                outAngles.append(angles[i])
                outVals.append(res[i])
            else:
                inAngles.append(angles[i])
                inVals.append(res[i])
        return [[inAngles, inVals], [outAngles, outVals]]

    def createVertsForPolyPath(self, x, y):
        verts = [[xi, yi] for xi, yi in zip(x, y)]
        verts.insert(0, [verts[0][0], -10e16])
        verts.append([verts[-1][0], -10e16])
        return verts


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

"""
***************************************************************************
    MeanAndStdDevPlot.py
    ---------------------
    Date                 : January 2013
    Copyright            : (C) 2013 by Victor Olaya
    Email                : volayaf at gmail dot com
***************************************************************************
*                                                                         *
*   This program is free software; you can redistribute it and/or modify  *
*   it under the terms of the GNU General Public License as published by  *
*   the Free Software Foundation; either version 2 of the License, or     *
*   (at your option) any later version.                                   *
*                                                                         *
***************************************************************************
"""

__author__ = 'Victor Olaya'
__date__ = 'January 2013'
__copyright__ = '(C) 2013, Victor Olaya'
# This will get replaced with a git SHA1 when you do a git archive
__revision__ = '$Format:%H$'

import matplotlib.pyplot as plt
import matplotlib.pylab as lab
import numpy as np
from PyQt4.QtCore import *
from qgis.core import *
from sextante.parameters.ParameterTable import ParameterTable
from sextante.parameters.ParameterTableField import ParameterTableField
from sextante.core.GeoAlgorithm import GeoAlgorithm
from sextante.outputs.OutputHTML import OutputHTML
from sextante.tools import *
from sextante.core.QGisLayers import QGisLayers

class MeanAndStdDevPlot(GeoAlgorithm):

    INPUT = "INPUT"
    OUTPUT = "OUTPUT"
    NAME_FIELD = "NAME_FIELD"
    MEAN_FIELD = "MEAN_FIELD"
    STDDEV_FIELD = "STDDEV_FIELD"

    def processAlgorithm(self, progress):
        uri = self.getParameterValue(self.INPUT)
        layer = QGisLayers.getObjectFromUri(uri)
        namefieldname = self.getParameterValue(self.NAME_FIELD)
        meanfieldname = self.getParameterValue(self.MEAN_FIELD)
        stddevfieldname = self.getParameterValue(self.STDDEV_FIELD)
        output = self.getOutputValue(self.OUTPUT)
        values = vector.getAttributeValues(layer, namefieldname, meanfieldname, stddevfieldname)
        plt.close()

        ind = np.arange(len(values[namefieldname]))
        width = 0.8
        plt.bar(ind, values[meanfieldname], width,
                    color='r',
                    yerr=values[stddevfieldname],
                    error_kw=dict(ecolor='yellow'))

        plt.xticks(ind, values[namefieldname], rotation = 45)
        plotFilename = output +".png"
        lab.savefig(plotFilename)
        f = open(output, "w")
        f.write("<img src=\"" + plotFilename + "\"/>")
        f.close()

    def defineCharacteristics(self):
        self.name = "Mean and standard deviation plot"
        self.group = "Graphics"
        self.addParameter(ParameterTable(self.INPUT, "Input table"))
        self.addParameter(ParameterTableField(self.NAME_FIELD, "Category name field", self.INPUT,ParameterTableField.DATA_TYPE_ANY))
        self.addParameter(ParameterTableField(self.MEAN_FIELD, "Mean field", self.INPUT))
        self.addParameter(ParameterTableField(self.STDDEV_FIELD, "StdDev field", self.INPUT))
        self.addOutput(OutputHTML(self.OUTPUT, "Output"))


"""
=================================================
Compare evoked responses for different conditions
=================================================

In this example, an Epochs object for visual and
auditory responses is created. Both conditions
are then accessed by their respective names to
create a sensor layout plot of the related
evoked responses.

"""

# Authors: Denis Engemann <denis.engemann@gmail.com>
#          Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>

# License: BSD (3-clause)

import matplotlib.pyplot as plt
import mne

from mne.viz import plot_evoked_topo
from mne.datasets import sample

print(__doc__)

data_path = sample.data_path()

###############################################################################
# Set parameters
raw_fname = data_path + '/MEG/sample/sample_audvis_filt-0-40_raw.fif'
event_fname = data_path + '/MEG/sample/sample_audvis_filt-0-40_raw-eve.fif'
event_id = 1
tmin = -0.2
tmax = 0.5

#   Setup for reading the raw data
raw = mne.io.read_raw_fif(raw_fname)
events = mne.read_events(event_fname)

#   Set up pick list: MEG + STI 014 - bad channels (modify to your needs)
include = []  # or stim channels ['STI 014']
# bad channels in raw.info['bads'] will be automatically excluded

#   Set up amplitude-peak rejection values for MEG channels
reject = dict(grad=4000e-13, mag=4e-12)

# pick MEG channels
picks = mne.pick_types(raw.info, meg=True, eeg=False, stim=False, eog=True,
                       include=include, exclude='bads')

# Create epochs including different events
event_id = {'audio/left': 1, 'audio/right': 2,
            'visual/left': 3, 'visual/right': 4}
epochs = mne.Epochs(raw, events, event_id, tmin, tmax,
                    picks=picks, baseline=(None, 0), reject=reject)

# Generate list of evoked objects from conditions names
evokeds = [epochs[name].average() for name in ('left', 'right')]

###############################################################################
# Show topography for two different conditions

colors = 'yellow', 'green'
title = 'MNE sample data - left vs right (A/V combined)'

plot_evoked_topo(evokeds, color=colors, title=title)

plt.show()


# -*- coding: utf-8 -*-
"""
Created on Wed Jan 13 23:21:29 2016
Defines a set of utility functions to be used for prediction
@author: Rupak Chakraborty
"""
import math
from trip import Trip
from gift import Gift
import random
import time
import pandas as pd
import operator

RADIUS_EARTH = 6773
NORTH_POLE_LAT = 90
NORTH_POLE_LONG = 0
EMPTY_SLEIGH_WEIGHT = 10
SLEIGH_CAPACITY = 1000
random.seed(time.time())
gift_filename = "Santa's Stolen Sleigh/gifts.csv"

"""
Calculates the haversine distance between two given points
The two points are the values of the latitude and longitude
in degrees

Params:
--------
lat_first - Latitude of the first point
long_first - Longitude of the first point
lat_second - Latitude of second point
long_second - Longitude of the second point

Returns:
---------
The haversine distance between the two given points i.e. a float

"""
def haversineDistance(lat_first,long_first,lat_second,long_second):

    lat_first = math.radians(lat_first)
    long_first = math.radians(long_first)
    lat_second = math.radians(lat_second)
    long_second = math.radians(long_second)

    sine_squared_lat = math.pow(math.sin((lat_first-lat_second)/2.0),2.0)
    sine_squared_long = math.pow(math.sin((long_first-long_second)/2.0),2.0)
    cos_lat_term = math.cos(lat_first)*math.cos(lat_second)*sine_squared_long
    total_term = cos_lat_term + sine_squared_lat

    distance = 2*RADIUS_EARTH*math.asin(math.sqrt(total_term))

    return distance

"""
Defines the fitness function for the trip list i.e. all deliveries
The total fitness is defined as the weighted sum of distances

Params:
--------
trip_list: A List of trips which Santa needs to take (A list containing the trip object)

Returns:
---------
Total Cost of the given trip list (i.e. Fitness)
"""
def tripFitness(trip_list):

    total_cost = 0

    for trip in trip_list:
        total_cost = total_cost + trip.trip_cost

    return total_cost

"""
Given a list of gifts calculates the cost of the trip (i.e. Weighted Distance)

Params:
--------
gift_list: A list of gifts in the order in which they have to be delivered

Returns:
---------

Cost of the trip with the given order of gifts (i.e. A Floating point number)
"""
def tripCost(gift_list):

    gift_size = len(gift_list)
    initial_gift_weight = tripWeightUtil(gift_list,0,gift_size-1)
    weighted_distance = initial_gift_weight*haversineDistance(NORTH_POLE_LAT,NORTH_POLE_LONG,gift_list[0].latitude,gift_list[0].longitude)

    for i in range(gift_size-1):
        remaining_weight = tripWeightUtil(gift_list,i+1,gift_size-1)
        distance = haversineDistance(gift_list[i].latitude,gift_list[i].longitude,gift_list[i+1].latitude,gift_list[i+1].longitude)
        weighted_distance = weighted_distance + remaining_weight*distance

    returning_distance = haversineDistance(gift_list[gift_size-1].latitude,gift_list[gift_size-1].longitude,NORTH_POLE_LAT,NORTH_POLE_LONG)
    weighted_distance = weighted_distance + EMPTY_SLEIGH_WEIGHT*returning_distance

    return weighted_distance

"""
Utility function to calculate the cumulative weight of gifts in a given range
Both ends of the range are included

Params:
--------

gift_list : List of gift objects
start_index : Starting index for gift list
end_index : Ending index of the gift list

Returns:
---------

Returns the sum of weights in a given range

"""
def tripWeightUtil(gift_list,start_index,end_index):

    total_weight = 0

    while start_index <= end_index:
        total_weight = total_weight + gift_list[start_index].weight
        start_index = start_index + 1

    return total_weight

"""
Applies the mutation operator on trip list i.e. swaps two trips

Params:
-------
trip_list: List containing the trips taken by Santa

Returns:
--------
A new list containing the trip list with values swapped
"""

def mutateTripList(trip_list):

    i,j = generateSwapIndices(len(trip_list))
    temp = trip_list[i]
    trip_list[i] = trip_list[j]
    trip_list[j] = temp

    return trip_list

"""
Applies the mutation operator on the gift list i.e. swaps two gifts in a list

Params:
-------
gift_list: List containing the gifts taken by Santa

Returns:
--------
A new list containing the gift list with values swapped
"""

def mutateGiftList(gift_list):

    i,j = generateSwapIndices(len(gift_list))
    temp = gift_list[i]
    gift_list[i] = gift_list[j]
    gift_list[j] = temp

    return gift_list

"""
Utility function to generate two distinct random integers from zero to a given range

Params:
--------
max_size: Integer containing the maximum limit for generation of the random integers

Returns:
--------
Two distinct random integers between 0 and a given max_size
"""
def generateSwapIndices(max_size):

    a = random.randint(0,max_size-1)
    b = random.randint(0,max_size-1)

    while b != a:
        b = random.randint(0,max_size)

    return a,b
"""
Returns the dataFrame containing the gift information

Params:
-------
String containing the filename from which the information is to be extracted

Returns:
--------

Pandas Dataframe object containing the gift information
"""
def getGiftList(filename):

    giftFrame = pd.read_csv(filename)
    gift_list = list([])

    for i in range(len(giftFrame)):

        gift_series = giftFrame.iloc[i]
        gift = Gift(gift_series.GiftId,gift_series.Latitude,gift_series.Longitude,gift_series.Weight)

        gift_list.append(gift)

    return gift_list;

"""
Sorts a given map by the values and returns a list containing th sorted tuples
"""
def sortMapByValues(map_to_sort):

    sorted_map = sorted(map_to_sort.items(), key=operator.itemgetter(1),reverse=False)
    return sorted_map

"""
Sorts the given population by its fitness value

Params:
-------
initial_population: List containing the initial population

Returns:
--------
List of tuples containing the indices of the initial population and its fitness
"""
def sortPopulationByFitness(initial_population):

    i = 0;
    fitness_population_map = {}

    for trip_gene in initial_population:
        fitness_population_map[i] = tripFitness(trip_gene)
        i = i + 1

    ordered_fitness_list = sortMapByValues(fitness_population_map)

    return ordered_fitness_list
"""
Given all the trips in a list returns the one with the maximum cost and its index

Params:
---------
trip_list: List of trips to be taken for delivery

Returns:
--------
The trip with the maximum cost and its corresponding index
"""
def maximumTripCost(trip_list):
    index = 0
    max_trip = trip_list[0]

    for i,trip in enumerate(trip_list):

        if trip.trip_cost > max_trip:
            max_trip = trip.trip_cost
            index = i

    return index,trip


"""
Agglomerative clustering with and without structure
===================================================

This example shows the effect of imposing a connectivity graph to capture
local structure in the data. The graph is simply the graph of 20 nearest
neighbors.

Two consequences of imposing a connectivity can be seen. First clustering
with a connectivity matrix is much faster.

Second, when using a connectivity matrix, average and complete linkage are
unstable and tend to create a few clusters that grow very quickly. Indeed,
average and complete linkage fight this percolation behavior by considering all
the distances between two clusters when merging them. The connectivity
graph breaks this mechanism. This effect is more pronounced for very
sparse graphs (try decreasing the number of neighbors in
kneighbors_graph) and with complete linkage. In particular, having a very
small number of neighbors in the graph, imposes a geometry that is
close to that of single linkage, which is well known to have this
percolation instability.
"""
# Authors: Gael Varoquaux, Nelle Varoquaux
# License: BSD 3 clause

import time
import matplotlib.pyplot as plt
import numpy as np

from sklearn.cluster import AgglomerativeClustering
from sklearn.neighbors import kneighbors_graph

# Generate sample data
n_samples = 1500
np.random.seed(0)
t = 1.5 * np.pi * (1 + 3 * np.random.rand(1, n_samples))
x = t * np.cos(t)
y = t * np.sin(t)

X = np.concatenate((x, y))
X += .7 * np.random.randn(2, n_samples)
X = X.T

# Create a graph capturing local connectivity. Larger number of neighbors
# will give more homogeneous clusters to the cost of computation
# time. A very large number of neighbors gives more evenly distributed
# cluster sizes, but may not impose the local manifold structure of
# the data
knn_graph = kneighbors_graph(X, 30, include_self=False)

for connectivity in (None, knn_graph):
    for n_clusters in (30, 3):
        plt.figure(figsize=(10, 4))
        for index, linkage in enumerate(('average', 'complete', 'ward')):
            plt.subplot(1, 3, index + 1)
            model = AgglomerativeClustering(linkage=linkage,
                                            connectivity=connectivity,
                                            n_clusters=n_clusters)
            t0 = time.time()
            model.fit(X)
            elapsed_time = time.time() - t0
            plt.scatter(X[:, 0], X[:, 1], c=model.labels_,
                        cmap=plt.cm.spectral)
            plt.title('linkage=%s (time %.2fs)' % (linkage, elapsed_time),
                      fontdict=dict(verticalalignment='top'))
            plt.axis('equal')
            plt.axis('off')

            plt.subplots_adjust(bottom=0, top=.89, wspace=0,
                                left=0, right=1)
            plt.suptitle('n_cluster=%i, connectivity=%r' %
                         (n_clusters, connectivity is not None), size=17)

plt.show()


#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import unittest

from pyspark.sql.functions import udf, pandas_udf, PandasUDFType
from pyspark.sql.types import DoubleType, StructType, StructField, LongType
from pyspark.sql.utils import ParseException, PythonException
from pyspark.rdd import PythonEvalType
from pyspark.testing.sqlutils import ReusedSQLTestCase, have_pandas, have_pyarrow, \
    pandas_requirement_message, pyarrow_requirement_message
from pyspark.testing.utils import QuietTest

@unittest.skipIf(
    not have_pandas or not have_pyarrow,
    pandas_requirement_message or pyarrow_requirement_message)  # type: ignore[arg-type]
class PandasUDFTests(ReusedSQLTestCase):

    def test_pandas_udf_basic(self):
        udf = pandas_udf(lambda x: x, DoubleType())
        self.assertEqual(udf.returnType, DoubleType())
        self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, DoubleType(), PandasUDFType.SCALAR)
        self.assertEqual(udf.returnType, DoubleType())
        self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, 'double', PandasUDFType.SCALAR)
        self.assertEqual(udf.returnType, DoubleType())
        self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, StructType([StructField("v", DoubleType())]),
                         PandasUDFType.GROUPED_MAP)
        self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())]))
        self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, 'v double', PandasUDFType.GROUPED_MAP)
        self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())]))
        self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, 'v double',
                         functionType=PandasUDFType.GROUPED_MAP)
        self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())]))
        self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, returnType='v double',
                         functionType=PandasUDFType.GROUPED_MAP)
        self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())]))
        self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

    def test_pandas_udf_decorator(self):
        @pandas_udf(DoubleType())
        def foo(x):
            return x
        self.assertEqual(foo.returnType, DoubleType())
        self.assertEqual(foo.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

        @pandas_udf(returnType=DoubleType())
        def foo(x):
            return x
        self.assertEqual(foo.returnType, DoubleType())
        self.assertEqual(foo.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

        schema = StructType([StructField("v", DoubleType())])

        @pandas_udf(schema, PandasUDFType.GROUPED_MAP)
        def foo(x):
            return x
        self.assertEqual(foo.returnType, schema)
        self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        @pandas_udf('v double', PandasUDFType.GROUPED_MAP)
        def foo(x):
            return x
        self.assertEqual(foo.returnType, schema)
        self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        @pandas_udf(schema, functionType=PandasUDFType.GROUPED_MAP)
        def foo(x):
            return x
        self.assertEqual(foo.returnType, schema)
        self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        @pandas_udf(returnType='double', functionType=PandasUDFType.SCALAR)
        def foo(x):
            return x
        self.assertEqual(foo.returnType, DoubleType())
        self.assertEqual(foo.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

        @pandas_udf(returnType=schema, functionType=PandasUDFType.GROUPED_MAP)
        def foo(x):
            return x
        self.assertEqual(foo.returnType, schema)
        self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

    def test_udf_wrong_arg(self):
        with QuietTest(self.sc):
            with self.assertRaises(ParseException):
                @pandas_udf('blah')
                def foo(x):
                    return x
            with self.assertRaisesRegexp(ValueError, 'Invalid return type.*None'):
                @pandas_udf(functionType=PandasUDFType.SCALAR)
                def foo(x):
                    return x
            with self.assertRaisesRegexp(ValueError, 'Invalid function'):
                @pandas_udf('double', 100)
                def foo(x):
                    return x

            with self.assertRaisesRegexp(ValueError, '0-arg pandas_udfs.*not.*supported'):
                pandas_udf(lambda: 1, LongType(), PandasUDFType.SCALAR)
            with self.assertRaisesRegexp(ValueError, '0-arg pandas_udfs.*not.*supported'):
                @pandas_udf(LongType(), PandasUDFType.SCALAR)
                def zero_with_type():
                    return 1

            with self.assertRaisesRegexp(TypeError, 'Invalid return type'):
                @pandas_udf(returnType=PandasUDFType.GROUPED_MAP)
                def foo(df):
                    return df
            with self.assertRaisesRegexp(TypeError, 'Invalid return type'):
                @pandas_udf(returnType='double', functionType=PandasUDFType.GROUPED_MAP)
                def foo(df):
                    return df
            with self.assertRaisesRegexp(ValueError, 'Invalid function'):
                @pandas_udf(returnType='k int, v double', functionType=PandasUDFType.GROUPED_MAP)
                def foo(k, v, w):
                    return k

    def test_stopiteration_in_udf(self):
        def foo(x):
            raise StopIteration()

        def foofoo(x, y):
            raise StopIteration()

        exc_message = "Caught StopIteration thrown from user's code; failing the task"
        df = self.spark.range(0, 100)

        # plain udf (test for SPARK-23754)
        self.assertRaisesRegexp(
            PythonException,
            exc_message,
            df.withColumn('v', udf(foo)('id')).collect
        )

        # pandas scalar udf
        self.assertRaisesRegexp(
            PythonException,
            exc_message,
            df.withColumn(
                'v', pandas_udf(foo, 'double', PandasUDFType.SCALAR)('id')
            ).collect
        )

        # pandas grouped map
        self.assertRaisesRegexp(
            PythonException,
            exc_message,
            df.groupBy('id').apply(
                pandas_udf(foo, df.schema, PandasUDFType.GROUPED_MAP)
            ).collect
        )

        self.assertRaisesRegexp(
            PythonException,
            exc_message,
            df.groupBy('id').apply(
                pandas_udf(foofoo, df.schema, PandasUDFType.GROUPED_MAP)
            ).collect
        )

        # pandas grouped agg
        self.assertRaisesRegexp(
            PythonException,
            exc_message,
            df.groupBy('id').agg(
                pandas_udf(foo, 'double', PandasUDFType.GROUPED_AGG)('id')
            ).collect
        )

    def test_pandas_udf_detect_unsafe_type_conversion(self):
        import pandas as pd
        import numpy as np

        values = [1.0] * 3
        pdf = pd.DataFrame({'A': values})
        df = self.spark.createDataFrame(pdf).repartition(1)

        @pandas_udf(returnType="int")
        def udf(column):
            return pd.Series(np.linspace(0, 1, len(column)))

        # Since 0.11.0, PyArrow supports the feature to raise an error for unsafe cast.
        with self.sql_conf({
                "spark.sql.execution.pandas.convertToArrowArraySafely": True}):
            with self.assertRaisesRegexp(Exception,
                                         "Exception thrown when converting pandas.Series"):
                df.select(['A']).withColumn('udf', udf('A')).collect()

        # Disabling Arrow safe type check.
        with self.sql_conf({
                "spark.sql.execution.pandas.convertToArrowArraySafely": False}):
            df.select(['A']).withColumn('udf', udf('A')).collect()

    def test_pandas_udf_arrow_overflow(self):
        import pandas as pd

        df = self.spark.range(0, 1)

        @pandas_udf(returnType="byte")
        def udf(column):
            return pd.Series([128] * len(column))

        # When enabling safe type check, Arrow 0.11.0+ disallows overflow cast.
        with self.sql_conf({
                "spark.sql.execution.pandas.convertToArrowArraySafely": True}):
            with self.assertRaisesRegexp(Exception,
                                         "Exception thrown when converting pandas.Series"):
                df.withColumn('udf', udf('id')).collect()

        # Disabling safe type check, let Arrow do the cast anyway.
        with self.sql_conf({"spark.sql.execution.pandas.convertToArrowArraySafely": False}):
            df.withColumn('udf', udf('id')).collect()

if __name__ == "__main__":
    from pyspark.sql.tests.test_pandas_udf import *  # noqa: F401

    try:
        import xmlrunner  # type: ignore[import]
        testRunner = xmlrunner.XMLTestRunner(output='target/test-reports', verbosity=2)
    except ImportError:
        testRunner = None
    unittest.main(testRunner=testRunner, verbosity=2)