#!/usr/bin/env python

import os
import openmc
import NCrystal as NC
import numpy as np
import matplotlib.pyplot as plt
openmc.config['cross_sections'] = '/home/marquezj/work/data/endfb-viii.0-hdf5/cross_sections.xml'

def create_model(m_mod, T0):
    m_refl = openmc.Material()
    m_refl.add_nuclide('O16', 1.0)
    m_refl.set_density('g/cm3', 1.0)
    m_refl.temperature = T0 #K
    
    width = 10
    pl1 = openmc.XPlane(x0=-width*0.25, boundary_type='vacuum')
    pl2 = openmc.XPlane(x0=+width*0.25)
    pl3 = openmc.XPlane(x0=+width*0.75, boundary_type='vacuum')
    
    cell01 = openmc.Cell(region=+pl1&-pl2, fill=m_mod)
    cell02 = openmc.Cell(region=+pl2&-pl3, fill=m_refl)
    
    g = openmc.Geometry(root=[cell01, cell02])
    
    t1 = openmc.Tally(name='moderator spectrum')
    f1 = openmc.EnergyFilter(np.geomspace(1e-5, 100, 200))
    f2 = openmc.CellFilter(cell01)
    t1.scores = ['flux']
    t1.filters = [f1, f2]
    t = openmc.Tallies([t1])
    
    s = openmc.Settings()
    s.run_mode = 'fixed source'
    s.particles = 10000
    s.batches = 100
    src = openmc.IndependentSource()
    src.energy = openmc.stats.Discrete(x=100.0, p=1.0)
    s.source = src
    
    return openmc.Model(geometry=g, settings=s, tallies=t)

def get_tally(spfile):
    tally = openmc.StatePoint(spfile).get_tally(name='moderator spectrum')
    df = tally.get_pandas_dataframe()
    E = df['energy high [eV]'].values
    phi = df['mean'].values
    return E, phi

def delete_if_exists(fn):
    try:
        os.remove(fn)
    except OSError:
        pass

if __name__ == '__main__':
    temperatures = np.linspace(293, 293.5, 6)
    spectra = []
    for T in temperatures:
        print(f'{T:.1f} K')
        cfg1 = 'freegas::H2O/1gcm3'
        m_nc = openmc.Material.from_ncrystal(cfg1+f';temp={T:.1f}K')
        delete_if_exists('statepoint.100.h5')
        delete_if_exists('summary.h5')
        create_model(m_nc, T0=293.2).run(output=False)
        spectra.append(get_tally('statepoint.100.h5'))

    plt.figure()
    for (E, spectrum), T in zip(spectra, temperatures):
        plt.semilogx(E, spectrum, label=f'{T:.1f} K')
    plt.legend()
    plt.xlim(1e-5,200)
    plt.xlabel('Energy [eV]')
    plt.ylabel('phi')
    plt.show()