#!/usr/bin/env python3

# Install with: pip install gximagecomputing

import gximagecomputing
import numpy as np
import matplotlib.pyplot as plt

os.environ['OMP_NUM_THREADS']='8' # number of parallel threads

gxi = gximagecomputing.GXRadioImageComputing()

# loading model generated by pyAMPP
# Generate with:
# gxbox --time 2015-06-08T20:16:09 --coords -280 90 --hpc --box-dims 200 128 64 --box-res 0.7 --pad-frac 0.1 --data-dir /home/user/pyampp/download --gxmodel-dir /home/user/pyampp/gx_models

model, model_dt = gxi.load_model_hdf("../pyAMPP/b3d_data_20220131T042000_dim256x128x64.h5")
# model, model_dt = gxi.load_model_sav("./model.sav") for IDL model provided by GX Simulator
ebtel_c, ebtel_dt = gxi.load_ebtel('ebtel_ss.sav')

xc=170.0
yc=390.0
dx=2.0
dy=2.0
Nx=150
Ny=150

freqlist=[5.8, 6.2, 6.6, 7.0, 7.4, 7.8, 8.2, 8.6, 9.0, 9.4, 9.8, 10.2, 10.6, 11.0, 11.4, 11.8]

Tbase=1e6
nbase=1e8

Q0=0.0217
a=0.3
b=2.7

# creating selective heating table
SHtable = np.zeros((7, 7))
w = [1.0, 1.0, 1.0, 1.2, 1.4, 1.6, 2.0]
for j in range(SHtable.shape[0]):
    for i in range(SHtable.shape[1]):
        SHtable[i, j] = w[i] * w[j]

SHtable[6, 6] = 0.1

result = gxi.synth_model(model, model_dt, ebtel_c, ebtel_dt,\
                         freqlist, Nx, Ny, xc, yc, dx, dy, Tbase, nbase, Q0, a, b)

resultSH = gxi.synth_model(model, model_dt, ebtel_c, ebtel_dt,\
                         freqlist, Nx, Ny, xc, yc, dx, dy, Tbase, nbase, Q0, a, b,\
                         SHtable=SHtable)

for i in range(result["TI"].shape[-1]):
    fig, axes = plt.subplots(1, 4, figsize=(15, 3))
    for ax, im in zip(axes.flat, (result["TI"][:, :, i], result["TV"][:, :, i],\
                                  resultSH["TI"][:, :, i], resultSH["TV"][:, :, i])):
        ax.imshow(im, interpolation=None)
    fig.show()