""" Miguel Nacenta 2022
Tests and simulation for the StairCaseProcedure class
"""

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
import numpy as np
import warnings as wr
import StaircaseProcedure as sp

# Controls of this script
numTrials = 50
drawGraph = True

# Prepare drawing stuff
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
plt.ion()
plt.show()

# Create an object
mySP = sp.StaircaseProcedure([10, 200], 1, 0.5, -1) # example

# Run experiment



for i in xrange(1,numTrials):
    newIntensity = mySP.getNextParameterValue()
    testPoint = mySP.getAssumedTestPoint()
    print("--------------------------")
    print("Test number: "+str(i))
    print("Test point = "+str(testPoint))
    print("The new intensity is: "+str(newIntensity))
    print("Did you see the difference? (y/n)")
    answer = raw_input()
    if (answer == "y"):
        mySP.update(1)
    else:
        mySP.update(0)

    estimatedLambda = mySP.getActualLambda()
    estimatedLambdaNorm = mySP.getEstimatedNormalisedLambda()
    print("==========================")
    print("my estimated threshold intensity is: "+str(estimatedLambda[0]))
    print("my estimated slope is: "+str(estimatedLambda[1]))
    print("my estimated threshold intensity (Norm) is: "+str(estimatedLambdaNorm[0]))
    print("my estimated slope (Norm) is: "+str(estimatedLambdaNorm[1]))


    # Now draw the lambda if you can
    if drawGraph:
        [X,Y,Z] = mySP.getProbabilitiesOfLambda()

        plt.cla()
        ax.plot_wireframe(X,Y,Z, rstride=1, cstride=1)
        plt.draw()
        plt.pause(0.001)

    #print(type(myPsy._probLambda))
    #print(myPsy._probLambda.shape)
    #print(X.shape)
    #print(Y.shape)
    #print(Z.shape)