#! /usr/bin/env python 
## \file estcj.py
## \ingroup basic_gas_dyn
##
## \brief Equilibrium Shock Tube Conditions, Junior
##
## Since 1968, we have been using the ESTC code by Malcolm McIntosh
## to compute the conditions in the end of the reflected shock tubes
## T1--T5 and HEG.  There are a number of problems in using the ESTC
## code, including uncertainty in updating the chemistry coefficients.
##
## This program, ESTCJ, moves away from the old chemistry model
## by making use of the CEA code from the NASA Glenn Research Center.
##
## \author PA Jacobs
##         Institute of Aerodynamics and Flow Technology
##         The German Aerospace Center, Goettingen.
##
## \version 24-Dec-02: First code.
##

import sys, math
from cea_gas import Gas

# -------------------------------------------------------------------
# First, global data.

DEBUG_ESTCJ = 0
PRINT_STATUS = 1

# -------------------------------------------------------------------
# Second, utility functions.

def shock_ideal(s1, us, s2):
    """Computes post-shock conditions, assuming ideal gas properties
    and a shock-stationary frame.
    Input:
        s1 pre-shock gas state
        us speed of gas coming into shock
        s2 post-shock gas state
    Returns both the shock-reference speed and the lab speed."""
    #
    M1 = us / s1.a
    V1 = us
    k = s1.k
    R = s1.R
    C_v = s1.C_v
    #
    s2.rho = s1.rho * (k + 1.0) * M1 * M1 \
             / (2.0 + (k - 1.0) * M1 * M1)
    s2.p = s1.p * (2.0 * k * M1 * M1 - (k - 1.0)) / (k + 1.0)
    s2.T = s2.p / (R * s2.rho)
    s2.e = s2.T * C_v
    #
    V2 = s1.rho / s2.rho * V1
    ug = V1 - V2
    s2.a = s1.a * math.sqrt(s2.T / s1.T)
    #
    s2.R = s1.R
    s2.k = s1.k
    s2.C_v = s2.C_v
    #
    return (V2,ug)


def shock_real(s1, us, s2):
    """Computes post-shock conditions, using high-temperature gas properties
    and a shock-stationary frame.
    Input:
        s1 pre-shock gas state
        us speed of gas coming into shock
        s2 post-shock gas state
    Returns both the shock-reference speed and the lab speed."""
    #
    (u2,ug) = shock_ideal( s1, us, s2 )
    if DEBUG_ESTCJ:
        print 'shock_real(): post-shock condition assuming ideal gas'
        s2.write_state(sys.stdout)
        print '    u2: %g m/s, ug: %g m/s' % (u2,ug)
    #
    V1    = us
    s1.EOS('pT');  # We assume that p1 and T1 are correct
    s2.EOS('pT');  # and that s2 contains a fair initial guess.
    #
    p1    = s1.p
    e1    = s1.e
    rho1  = s1.rho
    temp1 = p1 + rho1 * V1 * V1;             # momentum
    H1    = e1 + p1 / rho1 + 0.5 * V1 * V1;  # total enthalpy
    #
    rho_delta = 1.0
    T_delta = 1.0
    rho_tol = 1.0e-3; # tolerance in kg/m^3
    T_tol = 0.25;  # tolerance in degrees K
    #
    s0 = Gas(s1.gasName); # temporary workspace
    #
    # Update the estimates using the Newton-Raphson method.
    #
    for count in range(20):
        rho_save = s2.rho
        T_save = s2.T
        p_save = s2.p
        #
        p2 = p_save
        T2 = T_save
        s2.EOS('pT')
        e2 = s2.e
        rho2 = rho_save
        r2r1 = rho2 / rho1
        f1 = temp1 - p2 - rho1 * rho1 * V1 * V1 / rho2
        f2 = H1 - e2 - p2 / rho2 - 0.5 * V1 * V1 / (r2r1 * r2r1)
        f1_save = f1
        f2_save = f2
        #
        # Use finite differences to compute the Jacobian.
        d_rho = rho_save * 0.01
        d_T = T_save * 0.01
        #
        rho2 = rho_save + d_rho
        T2 = T_save
        s0.rho = rho2
        s0.T = T2
        s0.EOS('rhoT')
        p2 = s0.p
        e2 = s0.e
        f1 = temp1 - p2 - rho1 * rho1 * V1 * V1 / rho2
        f2 = H1 - e2 - p2 / rho2 - 0.5 * V1 * V1 / (r2r1 * r2r1)
        A = (f1 - f1_save) / d_rho
        C = (f2 - f2_save) / d_rho
        #
        rho2 = rho_save
        T2 = T_save + d_T
        s0.rho = rho2
        s0.T = T2
        s0.EOS('rhoT')
        e2 = s0.e
        p2 = s0.p
        f1 = temp1 - p2 - rho1 * rho1 * V1 * V1 / rho2
        f2 = H1 - e2 - p2 / rho2 - 0.5 * V1 * V1 / (r2r1 * r2r1)
        B = (f1 - f1_save) / d_T
        D = (f2 - f2_save) / d_T
        #
        # Invert Jacobian and multiply.
        det = A * D - B * C
        rho_delta = (D * f1_save - B * f2_save) / det
        T_delta = (-C * f1_save + A * f2_save) / det
        #
        rho_new = rho_save - rho_delta
        T_new   = T_save - T_delta
        #
        s2.rho = rho_new
        s2.T   = T_new
        s2.EOS('rhoT')
        #
        # Check convergence.
        if abs(rho_delta) < rho_tol and abs(T_delta) < T_tol: break
    #
    if DEBUG_ESTCJ:
        print ('shock_real(): count = %d, drho=%e, dT=%e' %
               (count, rho_delta, T_delta) )
    #
    # Back-out velocities via continuity.
    V2 = V1 * s1.rho / s2.rho
    ug = V1 - V2
    return (V2,ug)


def reflected_shock(s2, ug, s5):
    """Computes state5 which has brought the gas to rest at the
    end of the shock tube.
    Input:
        state2 is the post-incident-shock gas state
        ug     is the lab-frame velocity of the gas in state 2
        state5 is the stagnation state that will be filled in as
                  a side effect of this function
    Returns ur, the reflected shock speed."""
    #
    # As an initial guess, assume that we have a very strong shock
    # in an ideal gas.
    density_ratio = (s2.k + 1.0)/(s2.k - 1.0)
    ur_a = ug / density_ratio;
    (u5,ujunk) = shock_real( s2, (ur_a + ug), s5 )
    # The objective function is the difference in speeds,
    # units are m/s.  A value of zero for this function means
    # that, as the shock propagates upstream with speed ur,
    # the processed test gas is left in the end of the tube
    # with a velocity of zero in the laboratory frame.
    f_a = u5 - ur_a
    if DEBUG_ESTCJ:
        print 'Reflected shock: ur_a: %g, u5: %g' % (ur_a, u5)
    #
    # Now, we need to update this guess...use a secant update.
    #
    ur_b = 1.1 * ur_a
    (u5,ujunk) = shock_real( s2, (ur_b + ug), s5 )
    f_b = u5 - ur_b
    if DEBUG_ESTCJ:
        print 'Reflected shock: ur_b: %g, u5: %g' % (ur_b, u5)
    if abs(f_a) < abs(f_b):
        (f_a, f_b) = (f_b, f_a); (ur_a, ur_b) = (ur_b, ur_a)
    count = 0
    while abs(f_b) > 0.5 and count < 20:
        slope = (f_b - f_a) / (ur_b - ur_a)
        ur_c = ur_b - f_b / slope
        (u5,ujunk) = shock_real( s2, (ur_c + ug), s5 )
        f_c = u5 - ur_c
        if abs(f_c) < abs(f_b):
            ur_b = ur_c; f_b = f_c
        else:
            ur_a = ur_c; f_a = f_c
        count = count + 1
    #
    # At this point, ur_b should be out best guess.
    # Update the gas state data and return the best-guess value.
    #
    if count >= 20:
        print 'Reflected shock iteration did not converge.'
    (u5,ujunk) = shock_real( s2, (ur_b + ug), s5 )
    return ur_b


# -------------------------------------------------------------------
# Main script

if len(sys.argv) != 7:
    print 'Equilibrium Shock Tube Conditions'
    print 'Usage: estcj.pj <gasName> <p1> <T1> <us> <pe> <outFileName>'
    print 'Where: <gasName> is one of air, n2, co2, h2ne,'
    print '       <p1> is the shock tube fill pressure in Pa,'
    print '       <T1> is the shock tube fill temperature in K,'
    print '       <us> is the incident shock speed in m/s,'
    print '       <pe> is the equilibrium pressure in Pa,'
    print '       <outFileName> is obvious, hopefully.'
    sys.exit()

gasName = sys.argv[1]
p1 = float(sys.argv[2])
T1 = float(sys.argv[3])
us = float(sys.argv[4])
pe = float(sys.argv[5])
outFileName = sys.argv[6]

if DEBUG_ESTCJ: print 'estcj:', gasName, p1, T1, us, outFileName

fout = open(outFileName, 'w')
fout.write('estcj: Equilibrium Shock Tube Conditions\n')

fout.write('Input parameters:\n')
fout.write('    Gas is %s, p1: %g Pa, T1: %g K, us: %g m/s\n' %
           (gasName,p1,T1,us) )

if PRINT_STATUS: print 'Write pre-shock condition.'
fout.write('State 1: pre-shock condition\n')
state1 = Gas(gasName)
state1.set_from_pAndT(p1, T1)
state1.write_state(fout)
H1 = state1.e + state1.p/state1.rho

if PRINT_STATUS: print 'Start incident-shock calculation.'
state2 = Gas(gasName)
(u2,ug) = shock_real( state1, us, state2 )
fout.write('State 2: post-shock condition.\n')
state2.write_state(fout)
fout.write('    u2: %g m/s, ug: %g m/s\n' % (u2,ug) )

if PRINT_STATUS: print 'Start reflected-shock calculation.'
state5 = Gas(gasName)
ur = reflected_shock(state2, ug, state5)
fout.write('State 5: reflected-shock condition.\n')
state5.write_state(fout)
fout.write('    ur: %g m/s\n' % (ur,) )

if PRINT_STATUS: print 'Start calculation of isentropic relaxation.'
state5s = Gas(gasName)
state5s.set_from_pAndT(state5.p, state5.T);  # entropy is set
state5s.p = pe;                    # then pressure is relaxed
state5s.EOS('ps');                 # via an isentropic process
fout.write('State 5s: equilibrium condition (relaxation to pe)\n')
state5s.write_state(fout)
H5s = state5s.e + state5s.p/state5s.rho
fout.write('Enthalpy difference (H5s - H1): %g J/kg\n' % ((H5s - H1),) )

if PRINT_STATUS: print 'Done.'
fout.close()


# ------------------- end -----------------------
