# This script generate simulation data for the PI corrector + matrix multiplication.
import numpy as np


#############################################################################
# Response matrix, ref orbit and data input model
#############################################################################

# The inverted response matrix, we use an actual true dataset
respmat = np.load("respmat.npy")
# write to file
with open("respmat.txt", "w") as fp:
    for r in respmat.flatten():
        fp.write("{}\n".format(r))

# The orbit reference, again this is from true data
refmat = np.load("reforbit.npy")
# write to file
with open("reforbit.txt", "w") as fp:
    for r in refmat:
        fp.write("{}\n".format(r))

# BPM data generation
def generate_bpm_data(size, DC=0, A=0, F=100, N=0):
    """ Generate dummy BPM data, based on a simple model of DC + A.sin(F) + N.noise """

    sig = DC + A*np.sin(2*np.pi*F*np.arange(size)*1e-4+np.random.rand()*2*np.pi) + np.random.default_rng().normal(0, N, size)

    return sig.astype("int32")

#############################################################################
# Parameters
#############################################################################
N_BPM=respmat.shape[1]*2
N_PSC=respmat.shape[0]
size=200

N_MM_RND=23

#############################################################################
# Generate data input
#############################################################################
bpmdata=np.empty((N_BPM, size), dtype="int32")
for i, DC, A, F, N in zip(
        np.arange(N_BPM, dtype="int"),
        np.random.default_rng().uniform(-1e6, 1e6, N_BPM),
        np.random.default_rng().uniform(0, 200e3, N_BPM),
        np.random.default_rng().uniform(2,2000, N_BPM),
        np.random.default_rng().uniform(0, 20e3, N_BPM),):
    bpmdata[i] = generate_bpm_data(size, DC, A, F, N)

# write to file
# Warning, write x then y for each bpm, not all x then all y as the array is packed.
with open("bpmdata.txt", "w") as fp:
    for i in range(size):
        for j in range(N_BPM//2):
            fp.write("{} {} ".format(bpmdata[j,i], bpmdata[j+N_BPM//2,i]))
        fp.write("\n")    
        #fp.write(" ".join(bpmdata[:,i].astype("str"))+"\n")

#############################################################################
# Compute correction
#############################################################################
# change respmat shape, tiled
trespmat = np.zeros((N_PSC, N_BPM), dtype="int64")
trespmat[:50,:122] = respmat[:50]
trespmat[50:,122:] = respmat[50:]

K_A = 218
K_B = -186
K_iC = 325
K_D = -3225
C_N_RND = 20

## -----------------------
# Model computation

# Step 1: error computation
oe=bpmdata-refmat[:,np.newaxis]

# Step 2: matrix multiplication
mm=np.empty((N_PSC, size), dtype="int32")
for i in range(size):
    # Stage 1 : matmul
    _mm= np.matmul(oe[:,i].astype("int64"), trespmat.T, dtype='int64')
    mm[:,i] = _mm>>N_MM_RND

# Step 3: correction computation
corr = np.zeros((N_PSC, size), dtype="int64")

for i in range(1,size):
    _corr = ((mm[:,i]*K_A + mm[:,i-1]*K_B + corr[:,i-1]*K_D)*K_iC)
    corr[:,i] = np.clip(np.round(_corr/2**C_N_RND), -2**16, 2**16-1)



#############################################################################
# Write partial and results to files
#############################################################################
with open("orbiterror.txt", "w") as fp:
    for i in range(size):
        fp.write(" ".join(oe[:,i].astype("str"))+"\n")

with open("matmult.txt", "w") as fp:
    for i in range(size):
        fp.write(" ".join(mm[:,i].astype("str"))+"\n")

with open("corrout.txt", "w") as fp:
    for i in range(size):
        fp.write(" ".join(corr[:,i].astype("str"))+"\n")