import numpy
import random
import matplotlib
matplotlib.use('Cairo')
import matplotlib.pyplot as plt
from episode import Episode
from parameters import Parameters
from utils import ColorSpace
import cPickle
from explorer import Explorer

def main(arg):
    print 'arg = {0}'.format(arg)
    random.seed(1234)
    numberOfPoses = 5
    baseRadius = float(arg)
    radiusString = ''.join(`baseRadius`.split('.')) #for generating filename of plot
    print 'radiusString = {0}'.format(radiusString)
    a = range(-100,100)
    a.remove(0)
    radii = baseRadius * numpy.array(a, 'float')
    errorPerRadius = {}
    densityPerRadius = {}
    fileName = 'rolloutdump.pkl'
    f = open(fileName, 'r')
    data = cPickle.load(f)
    f.close()
    for radius in radii:
        errorPerRadius[radius] = {'raw':[]}
        densityPerRadius[radius] = {'raw':[]}
    #For parameterset theta
    for i in range(0,1):
        print i
        parameters = getParameters()
        observation, action = getRolloutData(data)
        ep = Episode(parameters, numberOfPoses, getJointNames())
        ep.explorer = Explorer()
        ep.setExploreJoint()
        #Calc density
        density = ep.getDensity(observation, action, parameters.getFullParameters(), ep.getSigmas(parameters))
        print 'density'
        print density
        #Calc gradient
        gradient, gradientlog = ep.estimateGradient(observation, action, parameters)
        print gradient
        #For various radii r
        for radius in radii:
            #Calc average density on hypersphere with radius r around theta
            print '     {0}'.format(radius)
            step = getStep(radius, len(gradient))
            estimatedDensity = density + numpy.dot(step, gradient)
            print 'estimated density'
            print estimatedDensity
            paramcp = parameters.getCopy()
            paramcp.setParameters(paramcp.getParameters() + step)
            print 'diff size in parameters = {0}'.format(numpy.linalg.norm(parameters.getParameters() - paramcp.getParameters()))
            realDensity = ep.getDensity(observation, action, paramcp.getFullParameters(), ep.getSigmas(parameters))
            print 'realDensity'
            print realDensity                
            error = abs(estimatedDensity - realDensity) / realDensity
            if realDensity == 0:
                errorPerRadius[radius]['raw'].append(0)
            densityPerRadius[radius]['raw'].append(realDensity)

    mean = []
    std = []
    for radius in radii:
        raw = numpy.array(errorPerRadius[radius]['raw'])
        mean.append(numpy.mean(raw))
        print 'raw = {0}'.format(raw)
        std.append(numpy.std(raw))

    plt.xlabel('Distance from initial parameters')
    plt.ylabel('(Error / density)')
    plt.errorbar(radii, mean, yerr=std, fmt='o')
    plt.savefig('errorwrtradius{0}.png'.format(radiusString), dpi = 250.0)    
    plt.clf()
    
    mean = []
    std = []
    for radius in radii:
        raw = numpy.array(densityPerRadius[radius]['raw'])
        mean.append(numpy.mean(raw))
        std.append(numpy.std(raw))
    print 'mean:'
    print mean
    print 'end'
    
    plt.xlabel('Distance from initial parameters')
    plt.ylabel('Density')
    plt.errorbar(radii, mean, yerr=std, fmt='o')
    plt.savefig('densitywrtradius{0}.png'.format(radiusString), dpi = 100.0)    
    plt.clf()

    f = open('temp.log', 'w')
    for i in range(0, len(mean)):
        line = '{0}, {1}, {2}\n'.format(radii[i], mean[i], std[i])
        f.write(line)
        #print line
    f.close()

def getJointNames():
    return ['laj1','laj2','laj3','laj4','llj2','llj3','llj4','llj5','llj6','raj1','raj2','raj3','raj4','rlj2','rlj3','rlj4','rlj5','rlj6']
        
def getRolloutData(data):
    timeSteps = data.keys()
    random.shuffle(timeSteps)
    time = timeSteps[0]
    x = data[time]['state']
    u = data[time]['action']
    return x,u
    
def getStep(radius, dimensions):
    source = []
    random.seed(1234)
    for i in range(0, dimensions):
        if i not in range(1,6):
            newElement = 1.0 #-1.0 + 2.0 * random.random()
        else:
            newElement = 0.0
        source.append(newElement)
    #for i in range(1, 6):
    #    source[i] = 0.0
    step = numpy.array(source)
    print 'step = {0}'.format(step)
    return radius * step / numpy.linalg.norm(step)
    
def getParameters():
    param = Parameters([0.5,
        0.83290912293510433, 0.83290912293510466, 0.78845736036427083, 0.83290912293510488, 0.83290912293510488, 
        -120.0, 19.0, -80.0, -80.0, 0.0,  37.0, -56.0, 27.0, 0.0,   -120.0, -20.0, 79.0, 79.0, 0.0,  37.0, -56.0, 27.0,  0.0, 
        -120.0, 19.0, -80.0, -80.0, 22.0, 37.0, -56.0, 25.0, -23.0, -120.0, -20.0, 79.0, 79.0, 22.0, 37.0, -56.0, 25.0, -23.0, 
        -78.0,  19.0, -80.0, -80.0, 23.0, 64.0, -74.0, 26.0, -25.0, -120.0, -53.0, 79.0, 15.0, 20.0, 60.0, -56.0, 17.0, -23.0, 
        -120.0, 19.0, -80.0, -80.0, 25.0, 61.0, -36.0, 0.0, -26.0, -120.0, -53.0, 79.0, 79.0, 23.0, 37.0, -59.0, 26.0, -24.0, 
        -120.0, 19.0, -80.0, -80.0, 0.0,  37.0, -71.0, 32.0, -1.0,  -120.0, -20.0, 79.0, 79.0, 0.0,  37.0, -71.0, 32.0, -1.0])
    return param

if __name__ == "__main__":
    import sys
    main(sys.argv[1])

f469c0362e41f49e174510077dd4af74a0121f28
import numpy
import random
import matplotlib
matplotlib.use('Cairo')
import matplotlib.pyplot as plt
from episode import Episode
from parameters import Parameters
from utils import ColorSpace
import cPickle
from explorer import Explorer

def main(arg):
    print 'arg = {0}'.format(arg)
    random.seed(1234)
    numberOfPoses = 5
    baseRadius = float(arg)
    radiusString = ''.join(`baseRadius`.split('.')) #for generating filename of plot
    print 'radiusString = {0}'.format(radiusString)
    a = range(-100,100)
    a.remove(0)
    radii = baseRadius * numpy.array(a, 'float')
    errorPerRadius = {}
    densityPerRadius = {}
    fileName = 'rolloutdump.pkl'
    f = open(fileName, 'r')
    data = cPickle.load(f)
    f.close()
    for radius in radii:
        errorPerRadius[radius] = {'raw':[]}
        densityPerRadius[radius] = {'raw':[]}
    #For parameterset theta
    for i in range(0,1):
        print i
        parameters = getParameters()
        observation, action = getRolloutData(data)
        ep = Episode(parameters, numberOfPoses, getJointNames())
        ep.explorer = Explorer()
        ep.setExploreJoint()
        #Calc density
        density = ep.getDensity(observation, action, parameters.getFullParameters(), ep.getSigmas(parameters))
        print 'density'
        print density
        #Calc gradient
        gradient, gradientlog = ep.estimateGradient(observation, action, parameters)
        print gradient
        #For various radii r
        for radius in radii:
            #Calc average density on hypersphere with radius r around theta
            print '     {0}'.format(radius)
            step = getStep(radius, len(gradient))
            estimatedDensity = density + numpy.dot(step, gradient)
            print 'estimated density'
            print estimatedDensity
            paramcp = parameters.getCopy()
            paramcp.setParameters(paramcp.getParameters() + step)
            print 'diff size in parameters = {0}'.format(numpy.linalg.norm(parameters.getParameters() - paramcp.getParameters()))
            realDensity = ep.getDensity(observation, action, paramcp.getFullParameters(), ep.getSigmas(parameters))
            print 'realDensity'
            print realDensity                
            error = abs(estimatedDensity - realDensity) / realDensity
            if realDensity == 0:
                errorPerRadius[radius]['raw'].append(0)
            densityPerRadius[radius]['raw'].append(realDensity)

    mean = []
    std = []
    for radius in radii:
        raw = numpy.array(errorPerRadius[radius]['raw'])
        mean.append(numpy.mean(raw))
        print 'raw = {0}'.format(raw)
        std.append(numpy.std(raw))

    plt.xlabel('Distance from initial parameters')
    plt.ylabel('(Error / density)')
    plt.errorbar(radii, mean, yerr=std, fmt='o')
    plt.savefig('errorwrtradius{0}.png'.format(radiusString), dpi = 250.0)    
    plt.clf()
    
    mean = []
    std = []
    for radius in radii:
        raw = numpy.array(densityPerRadius[radius]['raw'])
        mean.append(numpy.mean(raw))
        std.append(numpy.std(raw))
    print 'mean:'
    print mean
    print 'end'
    
    plt.xlabel('Distance from initial parameters')
    plt.ylabel('Density')
    plt.errorbar(radii, mean, yerr=std, fmt='o')
    plt.savefig('densitywrtradius{0}.png'.format(radiusString), dpi = 100.0)    
    plt.clf()

    f = open('temp.log', 'w')
    for i in range(0, len(mean)):
        line = '{0}, {1}, {2}\n'.format(radii[i], mean[i], std[i])
        f.write(line)
        #print line
    f.close()

def getJointNames():
    return ['laj1','laj2','laj3','laj4','llj2','llj3','llj4','llj5','llj6','raj1','raj2','raj3','raj4','rlj2','rlj3','rlj4','rlj5','rlj6']
        
def getRolloutData(data):
    timeSteps = data.keys()
    random.shuffle(timeSteps)
    time = timeSteps[0]
    x = data[time]['state']
    u = data[time]['action']
    return x,u
    
def getStep(radius, dimensions):
    source = []
    random.seed(1234)
    for i in range(0, dimensions):
        if i not in range(1,6):
            newElement = 1.0 #-1.0 + 2.0 * random.random()
        else:
            newElement = 0.0
        source.append(newElement)
    #for i in range(1, 6):
    #    source[i] = 0.0
    step = numpy.array(source)
    print 'step = {0}'.format(step)
    return radius * step / numpy.linalg.norm(step)
    
def getParameters():
    param = Parameters([0.5,
        0.83290912293510433, 0.83290912293510466, 0.78845736036427083, 0.83290912293510488, 0.83290912293510488, 
        -120.0, 19.0, -80.0, -80.0, 0.0,  37.0, -56.0, 27.0, 0.0,   -120.0, -20.0, 79.0, 79.0, 0.0,  37.0, -56.0, 27.0,  0.0, 
        -120.0, 19.0, -80.0, -80.0, 22.0, 37.0, -56.0, 25.0, -23.0, -120.0, -20.0, 79.0, 79.0, 22.0, 37.0, -56.0, 25.0, -23.0, 
        -78.0,  19.0, -80.0, -80.0, 23.0, 64.0, -74.0, 26.0, -25.0, -120.0, -53.0, 79.0, 15.0, 20.0, 60.0, -56.0, 17.0, -23.0, 
        -120.0, 19.0, -80.0, -80.0, 25.0, 61.0, -36.0, 0.0, -26.0, -120.0, -53.0, 79.0, 79.0, 23.0, 37.0, -59.0, 26.0, -24.0, 
        -120.0, 19.0, -80.0, -80.0, 0.0,  37.0, -71.0, 32.0, -1.0,  -120.0, -20.0, 79.0, 79.0, 0.0,  37.0, -71.0, 32.0, -1.0])
    return param

if __name__ == "__main__":
    import sys
    main(sys.argv[1])
