#!/usr/bin/env python

from embodied_ising import ising
import numpy as np
from sys import argv
from os import path, makedirs
from automatic_plot_helper import load_settings

import pickle

# --- COMPUTE MOTOR RESPONSE -------------------------------------------------------+

if len(argv) < 6:
    print("Usage: " + argv[0] + "<sim>" + "<moveCase>" + "<sensInd>" + "<sensor>" + "<gen>")
    exit(1)

if str(argv[2]) not in ("Move", "MoveOld", "MoveVelMot"):
    print("moveCase needs to be one of Move/MoveOld/MoveVelMot. Currently: " + str(argv[2]))
    exit(1)

loadfile = str(argv[1])
moveCase = str(argv[2])
sensInd = int(argv[3])
sensor = int(argv[4])
gen = int(argv[5])

def random_senses(sensor):
    # r_food = np.random.uniform(-180, 180)
    # org_sens = np.random.uniform(0, 1)
    # I.s[0] = r_food / 180
    # I.s[2] = np.tanh(org_sens) * 2 - 1

    # I.s[1] = np.random.uniform(-1, 1)
    # I.s[2] = np.random.uniform(-1, 1)

    for i in range(0, 3):  # 3 sensors
        if i != sensor:
            I.s[i] = np.random.uniform(-1, 1)

T = 10000
R = 5
Nsens = 100
mode = 'MonteCarlo'
sens = np.linspace(-1, 1, Nsens)

filename = 'save/' + loadfile + '/isings/gen[' + str(gen) + ']-isings.pickle'

startstr = 'Loading simulation:' + loadfile
print(startstr)

isings = pickle.load(open(filename, 'rb'))
settings = load_settings(loadfile)
numAgents = len(isings)

#  motor-dv, motor-dr, susc-motor-dv, susc-motor-dr
MV = np.zeros((R, numAgents))
MR = np.zeros((R, numAgents))
CMV = np.zeros((R, numAgents))
CMR = np.zeros((R, numAgents))

d = np.dtype([('MV', 'f8'), ('MR', 'f8'), ('CMV', 'f8'), ('CMR', 'f8')])
data = np.zeros((R, numAgents), dtype=d)

for rep in range(R):
    for agentNum, I in enumerate(isings):
        # print(agentNum)
        Mv = 0
        Mr = 0
        M2v = 0
        M2r = 0

        I.randomize_state()
        random_senses(sensor)
        I.s[sensor] = sens[sensInd]  # the sensor to hold constant

        for t in range(int(T / 10)):
            random_senses(sensor)
            if settings['ANN']:
                # update all other neurons a bunch of times
                for j in range(int(settings['thermalTime'])):
                    I.ANNStep()
            else:
                I.NoSensorGlauberStep()

        for t in range(T):
            random_senses(sensor)

            if settings['ANN']:
                # update all other neurons a bunch of times
                for j in range(int(settings['thermalTime'])):
                    I.ANNStep()
            else:
                I.NoSensorGlauberStep()

            # -1, 0, 1 for brake, nothing, accelerate
            # dv = (np.sum(I.s[-I.Msize1:]) / 2)
            # dr = (np.sum(I.s[-I.Msize:-I.Msize1]) / 2)
            if moveCase == 'Move':
                # this is the move mode where both motors are acceleration motors
                I.updateAcceleration()

                Mv += I.dv / float(T)  # dv
                Mr += I.ddr / float(T)  # dr

                M2v += I.dv ** 2 / float(T)
                M2r += I.ddr ** 2 / float(T)
            elif moveCase == 'MoveOld':
                # this is the move mode with one linear acceleration motor and one rotational velocity motor
                dv = (np.sum(I.s[-I.Msize1:]) / 2)
                dr = (np.sum(I.s[-I.Msize:-I.Msize1]) / 2)

                Mv += I.dv / float(T)  # dv
                Mr += I.dr / float(T)  # dr

                M2v += I.dv ** 2 / float(T)
                M2r += I.dr ** 2 / float(T)

            elif moveCase == 'MoveVelMot':
                # this is the move mode where both motors are velocity motors (same calculation as above)
                dv = (np.sum(I.s[-I.Msize1:]) / 2)
                dr = (np.sum(I.s[-I.Msize:-I.Msize1]) / 2)

                Mv += I.dv / float(T)  # dv
                Mr += I.dr / float(T)  # dr

                M2v += I.dv ** 2 / float(T)
                M2r += I.dr ** 2 / float(T)

        data['MV'][rep, agentNum] = Mv
        data['MR'][rep, agentNum] = Mr
        data['CMV'][rep, agentNum] = (M2v - Mv ** 2)
        data['CMR'][rep, agentNum] = (M2r - Mr**2)

folder = 'save/' + loadfile + '/Motor_Response/gen_' + str(gen) + '/'  + str(sensor) + '/'
file = 'MR-' + 'Nsens_' + str(Nsens) + \
       '-sensInd_' + str(sensInd) + \
       '-gen_' + str(gen) + \
       '.npy'

filename = folder + file

if not path.exists(folder):
    makedirs(folder)


np.save(filename, data)
savestr = 'Saving: ./.../' + file
print(savestr)
