'''
3D Heat Equation Based on Tran's model

Stefan: Extended to write SpaceEx file as output.

Stanley Bak, Stefan Schupp
April 2020
'''

import math

import numpy as np
from scipy.sparse import dia_matrix,find
from scipy.linalg import expm

import xml.etree.ElementTree as et
import xml.dom.minidom as minidom

import matplotlib.pyplot as plt

def main():
    '''main entry point'''

    # parameters
    samples_per_side = 5
    diffusity_const = 0.01
    heat_exchange_const = 0.5

    a_matrix = heat3d_dia(samples_per_side, diffusity_const, heat_exchange_const)

    print(f"A matrix has {a_matrix.shape[0]} dimensions")

    center_x = int(math.floor(samples_per_side/2.0))
    center_y = int(math.floor(samples_per_side/2.0))
    center_z = int(math.floor(samples_per_side/2.0))

    center_dim = center_z * samples_per_side**2 + center_y * samples_per_side + center_x

    print(f"Center Variable (Output): x{center_dim}")

    init_one = make_init(a_matrix, samples_per_side, 1.0)

    print(f"Non-zero Initial Variables:")

    for i, temp in enumerate(init_one):
        if temp > 1e-6:
            print(f"x{i} = ", end='')

    print("[0.9, 1.1]")

    pretty = True
    writeXML("heat_" + str(samples_per_side) + ".xml", a_matrix, pretty)
    writeCfg("heat_" + str(samples_per_side) + ".cfg",samples_per_side**3,init_one,center_dim)

    simulate = False

    if simulate:
        # plot several simulations
        init_temps = [0.9, 0.95, 1.0, 1.05, 1.1]
        step_size = 0.02
        num_steps = 2000

        s = samples_per_side
        filename = f'plot{s}.png'
        print(f"simulate = True, running simulations and saving to {filename}")

        mat = (a_matrix * step_size).toarray()
        print("computing matrix exponential")
        matrix_exponential = expm(mat)

        xs = [step_size * n for n in range(num_steps + 1)]
        ys = []

        temp_max = -np.inf
        time_max = -1

        for temp_index, init_temp in enumerate(init_temps):
            state = init_one * init_temp

            y = []
            ys.append(y)
            y.append(state[center_dim])

            for step in range(num_steps):
                print(f"sim {temp_index}/{len(init_temps)} - step {step}/{num_steps}")
                state = np.dot(matrix_exponential, state)
                temp = state[center_dim]
                y.append(temp)

                if temp > temp_max:
                    temp_max = temp
                    time_max = (1 + step) * step_size

        print(f"Maximum temp: {temp_max} at time {time_max}")

        # plot
        for y in ys:
            plt.plot(xs, y, 'k-')

        plt.title(f'Heat3D {s}x{s}x{s} ({s**3} dims)')
        plt.xlabel('Time')
        plt.ylabel(f'Center (x{center_dim}) Temperature')

        plt.savefig(filename)
        plt.show()

def make_init(a_matrix, samples, init_temp):
    '''returns an initial state vector'''

    vec = np.zeros((a_matrix.shape[0], 1), dtype=float)

    assert samples == 5 or (samples >= 10 and samples % 10 == 0), "init region isn't evenly divided by discretization"

    # maximum z point for initial region is 0.2 for 5 samples and 0.1 otherwise
    max_z = samples // 10 + 1 if samples >= 10 else 2 * samples // 10 + 1

    for z in range(max_z):
        zoffset = z * samples * samples

        for y in range(2 * samples // 10 + 1):
            yoffset = y * samples

            for x in range(4 * samples // 10 + 1):
                index = x + yoffset + zoffset

                vec[index] = init_temp

    return vec

def heat3d_dia(samples, diffusity_const, heat_exchange_const):
    'fast dia_matrix construction for heat3d dynamics'

    samples_sq = samples**2
    dims = samples**3
    step = 1.0 / (samples + 1)

    a = diffusity_const * 1.0 / step**2
    d = -2.0 * (a + a + a)

    data = np.zeros((7, dims))
    offsets = np.array([-samples_sq, -samples, -1, 0, 1, samples, samples_sq], dtype=float)

    # element with z = -1
    data[0, :-samples_sq] = a

    # element with y = -1
    for s in range(samples):
        start = s * samples_sq
        end = (s + 1) * (samples_sq) - samples
        data[1, start:end] = a

    # element with x = -1
    for s in range(samples_sq):
        start = s * samples
        end = (s + 1) * (samples) - 1
        data[2, start:end] = a

    #### diagonal element ####
    data[3, :] = d     # (prefill)

    # adjust when z = 0 or z = samples-1
    data[3, :samples_sq] += a
    data[3, -samples_sq:] += a

    # adjust when y = 0 or y = samples-1
    for z in range(samples):
        z_offset = z * samples_sq

        data[3, z_offset:z_offset + samples] += a
        data[3, z_offset + samples_sq - samples:z_offset + samples_sq] += a

    # adjust when x = 0 (and add diffusion term when x = samples-1)
    for z in range(samples):
        for y in range(samples):
            offset = z * samples_sq + y * samples

            data[3, offset] += a

            data[3, offset + samples - 1] += a / (1+heat_exchange_const * step)

    #### end diagnal element ####
    # element with x = +1
    for s in range(samples_sq):
        start = 1 + s * samples
        end = (s + 1) * samples
        data[4, start:end] = a

    # element with y = +1
    for s in range(samples):
        start = s * samples_sq + samples
        end = (s + 1) * (samples_sq)
        data[5, start:end] = a

    # element with z = +1
    data[6, samples_sq:] = a

    rv = dia_matrix((data, offsets), shape=(dims, dims))
    assert np.may_share_memory(rv.data, data)     # make sure we didn't copy memory

    return rv

def writeCfg(outfilename, numVars, init_one, centervar):
    ''' create SpaceEx config file '''
    fileobj = open(outfilename,'w')

    # some default values
    fileobj.write("system = \"system\"\n\
scenario = \"supp\"\n\
directions = \"uni8\"\n\
sampling-time = 0.002\n\
time-horizon = 40\n\
iter-max = -1\n\
output-variables = \"t,x_" + str(centervar) + "\"\n\
output-format = \"GEN\"\n\
rel-err = 1.0e-12\n\
abs-err = 1.0e-13\n")

    initialStates = []
    for i in range(len(init_one)):
        if init_one[i] == 0:
            initialStates.append("x_" + str(i) + " == 0")
        else:
            initialStates.append("0.9 <= x_" + str(i) + " <= 1.1")

    initialStates.append("t == 0")

    fileobj.write("initially = \"" + " & ".join(initialStates) + "\"")
    fileobj.close()

def writeXML(outfilename, matrix, pretty=False):
    ''' create SpaceEx model file '''

    #preamble/document structure
    root = et.Element('sspaceex')
    root.set('xmlns','http://www-verimag.imag.fr/xml-namespaces/sspaceex')
    root.set('version','0.2')
    root.set('math','SpaceEx')
    component = et.SubElement(root, 'component')
    component.set('id','plant_component_template')

    # binding component system
    system = et.SubElement(root,'component')
    system.set('id','system')

    # set up variables, also add to binding
    for rowIndex in range(matrix.shape[0]):
        varString = "x_" + str(rowIndex)
        var = et.SubElement(component,'param')
        var.set('name',varString)
        var.set('type','real')
        var.set('local','false')
        var.set('d1','1')
        var.set('d2','1')
        var.set('dynamics','any')

        # system component parameter bindings
        var_sys = et.SubElement(system,'param')
        var_sys.set('name',varString)
        var_sys.set('type','real')
        var_sys.set('local','false')
        var_sys.set('d1','1')
        var_sys.set('d2','1')
        var_sys.set('dynamics','any')
        var_sys.set('controlled','true')

    # add t for time
    var = et.SubElement(component,'param')
    var.set('name','t')
    var.set('type','real')
    var.set('local','false')
    var.set('d1','1')
    var.set('d2','1')
    var.set('dynamics','any')

    # system component parameter bindings
    var_sys = et.SubElement(system,'param')
    var_sys.set('name','t')
    var_sys.set('type','real')
    var_sys.set('local','false')
    var_sys.set('d1','1')
    var_sys.set('d2','1')
    var_sys.set('dynamics','any')
    var_sys.set('controlled','true')

    # component bindings
    plant_binding = et.SubElement(system, 'bind')
    plant_binding.set('component','plant_component_template')
    plant_binding.set('as','plant')

    # write mapping - this needs to be done AFTER specification of the params in the system component
    for rowIndex in range(matrix.shape[0]):
        varString = "x_" + str(rowIndex)
        # variable bindings
        mapping = et.SubElement(plant_binding,'map')
        mapping.set('key',varString)
        mapping.text = varString

    mapping = et.SubElement(plant_binding,'map')
    mapping.set('key','t')
    mapping.text = 't'

    location = et.SubElement(component, 'location')
    location.set('id','1')
    location.set('name','loc1')
    invariant = et.SubElement(location,'invariant')
    invariant.text = ''

    # set up dynamics
    dynamics = et.SubElement(location,'flow')
    dynamicslist = []
    for rowIndex in range(matrix.shape[0]):
        derivative = "x_" + str(rowIndex) + "' == "
        row = matrix.getrow(rowIndex)
        data = find(row)
        flowterms = []
        for nzCol in range(len(data[1])):
            coeffString = "x_" + str(data[1][nzCol])
            flowterms.append(str(data[2][nzCol]) + "*" + coeffString)

        dynamicslist.append(derivative + " + ".join(flowterms))

    dynamicslist.append("t' == 1")

    dynamics.text = " &\n".join(dynamicslist)

    if pretty:
        fileobj = open(outfilename, 'w')
        fileobj.write(prettyXML(root))
        fileobj.close()
    else:
        tree = et.ElementTree(root)
        tree.write(outfilename,encoding="unicode",xml_declaration=True,)

def prettyXML(elem):
    ''' adds linebreaks and indentation (not provided by ElemTree but via minidom) '''
    rough_string = et.tostring(elem, 'utf-8')
    reparsed = minidom.parseString(rough_string)
    return reparsed.toprettyxml(indent="\t")


if __name__ == '__main__':
    main()