#=##############################################################################
# DESCRIPTION
    Vahana vehicle geometry generation.

# AUTHORSHIP
  * Author    : Eduardo J. Alvarez
  * Email     : Edo.AlvarezR@gmail.com
  * Created   : Feb 2021
  * License   : MIT
=###############################################################################

#=

Geometry functions (e.g., `generategeometry_vahana`) are expected have the
following format:

    OUTPUTS
    * `system::vlm.WingSystem`:        System of all FLOWVLM objects.
    * `rotors::Array{vlm.Rotor, 1}`:   System of all FLOWVLM Rotor objects.
    * `tilting_systems::Tuple(vlm.WingSystem, ...)`:   Tuple of all FLOWVLM
                                        tilting objects, where
                                        `tilting_systems[i]` contains the i-th
                                        FLOWVLM system of lifting surfaces and
                                        rotors that tilt together.
    * `rotor_systems::Tuple(Array{vlm.Rotor,1}, ...)`:   Tuple of groups of
                                        Rotors that share a common RPM.
    * `vlm_system::vlm.WingSystem`:    System of all FLOWVLM objects to be
                                        solved through the VLM solver.
    * `wake_system::vlm.WingSystem`:   System of all FLOWVLM objects that will
                                        shed a VPM wake.
    * `fuselage::gt.Grid`:             Grid object of the fuselage.
    * `grounds::gt.Grid`:              Grid object of the ground.
    * `strn::String`                   String containing all vtk file names
                                        that will be generated if `system` is
                                        saved.

    NOTES
    * Maneuver functions are expected to return `(Vaircraft, angles, RPMs)`,
        where `Vaircraft(t)` is the translation velocity of both `system`
        and `fuselage`, `angles(t)[i]` is the tilt angle of
        `tilting_systems[i]`, `RPMs(t)[i]` is the RPM of
        `rotor_systems[i]`. `angles(t)[end]` is the tilt angle of the
        entire aircraft.
=#

"""
    Generates the geometry of Vahana aircraft
"""
function generate_geometry_vahana(;
                                    # AIRCRAFT OPTIONS
                                    tiltrotor_file="VahanaTilt.csv", # Tilt-rotor
                                    stackedrotor_file="VahanaStacked.csv", # Stacked-rotor
                                    data_path=uns.def_data_path,
                                    xfoil=true,                     # Run XFOIL
                                    n_factor::Int=1,                # Refinement factor
                                    add_rotors=true,
                                    VehicleType=uns.VLMVehicle,     # Type of vehicle to generate
                                    # OUTPUT OPTIONS
                                    run_name="vahana",
                                    verbose=true, v_lvl=0)

    ############################################################################
    # PARAMETERS
    ############################################################################

    if verbose; println("\t"^(v_lvl)*"Defining parameters..."); end;

    # ------------ GEOMETRIC PARAMETERS ------------------------------------
    init_ori = 90.0               # Initial orientation of wings
    init_ori_prop = 0*90.0        # Initial orientation of rotors
    init_ori_stackedprop = 30.0   # Initial orientation of stacked rotors

    # Rotors
    # NOTE: The radii defined here must match the radius in rotor_file
    R_w = 0.75                    # (m) main wing rotor radius
    R_tw = R_w                    # (m) tandem wing rotor radius
    CW_w = true                   # Clockwise main wing rotor
    CW_tw = false                 # Clockwise tandem wing rotor
    np_w = 2                      # Number of propellers on wing on each side
    np_tw = 2                     # Number of propellers on tandem wing each side
    n_ccb = 7*n_factor            # Number of CCBlade elements
    blade_r = 1/20                # Geometric expansion between blade elements
    ReD07 = 1.5e6                 # Prop Reynolds number of relative velocity at r/R=0.7
    ReD = ReD07/0.7               # Reynolds at tip
    pitch = 0.0                   # (deg) pitch of propellers
    soD = 0.1                     # Tip-to-tip distance of rotors over diameter
    xoc_offset_main = 0.175       # Axial distance of rotors from LE over chord
    xoc_offset_tandem = 0.10      # Axial distance of rotors from LE over chord
    main_outtilt = 10             # Mounted out-tilt of tandem-wing rotors
    tandem_pitchtilt = 10         # Mounted pitch-tilt of tandem-wing rotors

    # Stacked rotors
    stacked = [np_w, np_w+1]      # Index of wing rotors that are stacked rotors
    stckd_xoc_offset = -1.40      # Axial distance of rotors from LE over chord
    stckd_zoc_offset = 0.10       # Distance of stacked rotors above wing over chord
    stckd_zoR_offset = -0.05      # Stacking distance between stacked rotors over R
    stckd_corotating = true       # Co-rotating stacked rotors
    stckd_phase = -10             # Initial phase difference of stacked rotors
    stckd_pitch = 7.5             # Pitch difference of stacked rotors

    # Wing
    b_w = 5.86                    # (m) span
    AR_w = 7.4                    # Aspect ratio
    tr_w = 1.0                    # Taper ratio
    twist_r_w = 7.5               # (deg) twist at root
    twist_t_w = twist_r_w         # (deg) twist at tip
    lambda_w = main_outtilt       # (deg) sweep
    gamma_w = 5.0                 # (deg) dihedral
    n_w = 2*12*n_factor           # Number of horseshoes per side of wing
    r_w = 2.0                     # Horseshoe expansion ratio
    md_w = 0.9                    # Ratio of length of middle section
    pivot_w = 1/4                 # Pivot point along chord of tilt-wing

    # Wing winglets
    b_wl = b_w/4                  # (m) span of winglet from top to bottom
    AR_wl = 3.0                   # Aspect ratio
    tr_wl = (b_wl/AR_wl)/(b_w/AR_w/tr_w)      # Taper ratio
    twist_r_wl = 2.5              # (deg) twist at root
    twist_t_wl = 0.0              # (deg) twist at tip
    lambda_wl = 40.0              # (deg) sweep
    gamma_wl = 15.0               # (deg) dihedral
    n_wl = 2*4*n_factor           # Number of horseshoes per side of wing
    r_wl = 2.0                    # Horseshoe expansion ratio

    # Tandem wing
    b_tw = b_w*1.0                # (m) span
    AR_tw = 9.5                   # Aspect ratio
    tr_tw = 1.0                   # Taper ratio
    twist_r_tw = 4.0              # (deg) twist at root
    twist_t_tw = twist_r_tw       # (deg) twist at tip
    lambda_tw = 0.0               # (deg) sweep
    gamma_tw = 0.0                # (deg) dihedral
    n_tw = 2*n_w                  # Number of horseshoes per side of wing
    r_tw = r_w                    # Horseshoe expansion ratio
    md_tw = 0.2                   # Ratio of length of middle section
    pivot_tw = pivot_w            # Pivot point along chord of tilt-wing

    # Fuselage
    l_f = 5.86                    # (m) length
    h_f = 2.81*2/3                # (m) height
    c_pos_f1 = 0                  # (m) span position of first chord
    c_pos_f2 = 0.2*h_f            # (m) span position of second chord
    c_pos_f3 = 0.7*h_f
    c_pos_f4 = h_f
    c_f1 = 0.66*l_f               # (m) length of first chord
    c_f2 = 0.75*l_f
    c_f3 = 0.75*l_f
    c_f4 = 0.50*l_f
    x_pos_f1 = 0.025*l_f          # (m) x-position of first chord
    x_pos_f2 = 0.0*l_f
    x_pos_f3 = 0.20*l_f
    x_pos_f4 = 0.50*l_f

    # ------------ ASSEMBLY PARAMETERS ------------------------------------
    # Position of wings on fuselage
    h_pos_w = 0.90*h_f            # (m) height position of wing
    h_pos_tw = 0.15*h_f           # (m) height position of tandem wing
    l_pos_w = 0.95*l_f-b_w/AR_w/tr_w   # (m) length position of wing
    l_pos_tw = 0.05*l_f           # (m) length position of tandem wing

    # Position of propellers on main wing
    # d_prop_w = d_factor*b_w/2/np_w# Distance between propellers on wing
    d_prop_w = (1+soD)*(2*R_w)
    y_pos_prop_w = Float64[b_w/2 - i*d_prop_w for i in 0:np_w-1] # y-positions

    # Position of propellers on tandem wing
    # d_prop_tw = d_factor*b_tw/2/np_tw# Distance between propellers on wing
    d_prop_tw = (1+soD)*(2*R_tw)
    y_pos_prop_tw = Float64[b_tw/2 - i*d_prop_tw for i in 0:np_tw-1] # y-positions

    ############################################################################
    # ASSEMBLY
    ############################################################################

    if verbose; println("\t"^(v_lvl)*"Generating geometry..."); end;

    # ------------ ROTORS ------------------------------------------------
    # Generates base propellers (one on each rotation orientation)
    if add_rotors
        tiltrotors = vlm.Rotor[]
        if verbose; println("\t"^(v_lvl+1)*"Generating first tilt-rotor..."); end;
        push!(tiltrotors, uns.generate_rotor(tiltrotor_file; pitch=pitch,
                                                n=n_ccb, blade_r=blade_r, CW=!CW_w, ReD=ReD,
                                                verbose=verbose, v_lvl=v_lvl+2, xfoil=xfoil,
                                                data_path=data_path, plot_disc=false))
        if verbose; println("\t"^(v_lvl+1)*"Generating second tilt-rotor..."); end;
        # push!(tiltrotors, generate_rotor(pitch; n=n_ccb, blade_r=blade_r, CW=CW_w, ReD=ReD,
        #                         verbose=verbose, xfoil=xfoil, rotor_file=tiltrotor_file))
        push!(tiltrotors, vlm.Rotor(!tiltrotors[1].CW, tiltrotors[1].r,
                                  tiltrotors[1].chord, tiltrotors[1].theta,
                                  tiltrotors[1].LE_x, tiltrotors[1].LE_z,
                                  tiltrotors[1].B, tiltrotors[1].airfoils))
        vlm.initialize(tiltrotors[2], tiltrotors[1].m)

        stackedrotors = vlm.Rotor[]
        if verbose; println("\t"^(v_lvl+1)*"Generating first stacked-rotor..."); end;
        push!(stackedrotors, uns.generate_rotor(stackedrotor_file; pitch=pitch,
                                                n=n_ccb, blade_r=blade_r, CW=!CW_w, ReD=ReD,
                                                verbose=verbose, v_lvl=v_lvl+2, xfoil=xfoil,
                                                data_path=data_path, plot_disc=false))
        if verbose; println("\t"^(v_lvl+1)*"Generating second stacked-rotor..."); end;
        # push!(stackedrotors, generate_rotor(pitch; n=n_ccb, blade_r=blade_r, CW=CW_w, ReD=ReD,
        #                         verbose=verbose, xfoil=xfoil, rotor_file=stackedrotor_file))
        push!(stackedrotors, vlm.Rotor(!stackedrotors[1].CW, stackedrotors[1].r,
                                  stackedrotors[1].chord, stackedrotors[1].theta,
                                  stackedrotors[1].LE_x, stackedrotors[1].LE_z,
                                  stackedrotors[1].B, stackedrotors[1].airfoils))
        vlm.initialize(stackedrotors[2], stackedrotors[1].m)
    end

    # ------------ MAIN WING ---------------------------------------------
    # Generates wing
    if verbose; println("\t"^(v_lvl+1)*"Generating main wing assembly..."); end;

    # wing = vlm.simpleWing(b_w, AR_w, tr_w, twist_r_w, lambda_w, gamma_w;
    #                                           twist_tip=twist_t_w, n=n_w, r=r_w)

    # Middle section
    # tip = b/ar
    # root = tip/tr
    # newtip = root + md*(tip-root) = tip/tr + md*tip*(1 - 1/tr) = tip * ( md + (1/tr)*(1-md) )
    # newtr = newtip/root = 1 + md*(tip/root - 1) = 1 + md*(tr - 1)
    # newar = b/newtip = ar / ( md + (1/tr)*(1-md) )
    pos_md_w = [0.0, md_w]
    clen_md_w = [1/tr_w, md_w + (1/tr_w)*(1-md_w)]
    twist_md_w = [twist_r_w, twist_r_w + md_w*(twist_t_w-twist_r_w)]
    wing_md = vlm.complexWing(b_w, AR_w, ceil(Int, md_w*n_w), pos_md_w, clen_md_w, twist_md_w,
                          lambda_w*ones(1), gamma_w*ones(1);
                          symmetric=true, chordalign=0.0,
                          _ign1=true)
    # Left section
    # pos_l_w = [-(1-md_w), -0]
    pos_l_w = [0, -(1-md_w)]   # Here I'm defining this wing section from right to left
    clen_l_w = [1, clen_md_w[end]]
    twist_l_w = [twist_t_w, twist_md_w[end]]
    wing_L = vlm.complexWing(b_w, AR_w, ceil(Int, (1-md_w)*n_w/2), pos_l_w, clen_l_w,
                            twist_l_w, -lambda_w*ones(1), -gamma_w*ones(1);
                            symmetric=false, chordalign=0.0,
                            _ign1=true)
    # Right section
    pos_r_w = [0, (1-md_w)]
    clen_r_w = [clen_md_w[end], 1]
    twist_r_w = [twist_md_w[end], twist_t_w]
    wing_R = vlm.complexWing(b_w, AR_w, ceil(Int, (1-md_w)*n_w/2), pos_r_w, clen_r_w,
                            twist_r_w, lambda_w*ones(1), gamma_w*ones(1);
                            symmetric=false, chordalign=0.0,
                            _ign1=true)

    O_w_R = (md_w*b_w/2)*[0, 1, 0]
    O_w_L = [1 0 0; 0 -1 0; 0 0 1]*O_w_R
    vlm.setcoordsystem(wing_R, O_w_R, Float64[1 0 0; 0 1 0; 0 0 1])
    vlm.setcoordsystem(wing_L, O_w_L, Float64[1 0 0; 0 1 0; 0 0 1])

    # Generates winglets
    winglet_R = vlm.simpleWing(b_wl, AR_wl, tr_wl, twist_r_wl, lambda_wl, gamma_wl;
                                          twist_tip=twist_t_wl, n=n_wl, r=r_wl)
    winglet_L = vlm.simpleWing(b_wl, AR_wl, tr_wl, twist_r_wl, lambda_wl, gamma_wl;
                                          twist_tip=twist_t_wl, n=n_wl, r=r_wl)

    O_wl_R = (b_w/2)*[0, 1, 0]
    O_wl_R += ((1-md_w)*b_w/2)*[tan(lambda_w*pi/180), 0, tan(gamma_w*pi/180)]
    O_wl_L = [1 0 0; 0 -1 0; 0 0 1]*O_wl_R
    Oaxis_wl_R = gt.rotation_matrix(0.0, 0.0, 90.0)
    Oaxis_wl_L = gt.rotation_matrix(0.0, 0.0, -90.0)
    vlm.setcoordsystem(winglet_R, O_wl_R, Oaxis_wl_R)
    vlm.setcoordsystem(winglet_L, O_wl_L, Oaxis_wl_L)

    ## Generates propellers on wing (from right to left)
    if add_rotors
        if verbose; println("\t"^(v_lvl+2)*"Generating main wing propellers..."); end;
        O_prop_w = [
                     (ypos - md_w*b_w/2)*[tan(lambda_w*pi/180), 0, -tan(gamma_w*pi/180)] +
                     ypos*[0, 1, 0] +
                     [-(i in stacked ? stckd_xoc_offset : xoc_offset_main)*AR_w/b_w, 0, 0] +
                     [0, 0, (i in stacked ? stckd_zoc_offset*tan(gamma_w*pi/180) : 0)*AR_w/b_w]
                     for (i, ypos) in enumerate(y_pos_prop_w)]
        props_w = vlm.Rotor[]
        props_w_stacked_up = vlm.Rotor[]
        props_w_stacked_low = vlm.Rotor[]
        for i in 1:2*np_w
            right = i<=np_w    # Indicates which side of the wing

            copy_prop = i in stacked ? stackedrotors[1+i%2] : tiltrotors[1+i%2]
            this_prop = deepcopy(copy_prop) # Alternates rotation orientation

            this_O = O_prop_w[ right ? i : np_w-(i-np_w-1)]   # Chooses position
            this_O = [1 0 0; 0 (-1)^!right 0; 0 0 1]*this_O   # Places it in correct side
            this_Oaxis = i in stacked ? gt.rotation_matrix(0.0, 90, 0.0) :
                                        gt.rotation_matrix((-1)^(i%2==0)*main_outtilt, 0.0, 0.0)

            vlm.setcoordsystem(this_prop, this_O, this_Oaxis; user=true)

            # Rotates props to be tip to tip
            vlm.rotate(this_prop, (-1)^(!CW_w) * (i in stacked ? init_ori_stackedprop : init_ori_prop))

            # Adds the original polars that don't get copied in deepcopy
            this_prop.airfoils = copy_prop.airfoils
            this_prop._polars = copy_prop._polars
            this_prop._polarroot = copy_prop._polarroot
            this_prop._polartip = copy_prop._polartip

            if !(i in stacked)
                push!(props_w, this_prop)

            else
                push!(props_w_stacked_up, this_prop)

                copy_prop = deepcopy(stackedrotors[ 1+(i+1*!stckd_corotating)%2 ])

                # Generate lower rotor of stack and add pitch
                if stckd_pitch != 0
                    this_prop = vlm.Rotor(copy_prop.CW, copy_prop.r,
                                              copy_prop.chord, copy_prop.theta .+ stckd_pitch,
                                              copy_prop.LE_x, copy_prop.LE_z,
                                              copy_prop.B, copy_prop.airfoils)
                    vlm.initialize(this_prop, copy_prop.m)
                else
                    this_prop = copy_prop
                end

                this_O += R_w*[0, 0, stckd_zoR_offset]
                vlm.setcoordsystem(this_prop, this_O, this_Oaxis; user=true)

                # Rotates props to be tip to tip
                vlm.rotate(this_prop, (-1)^(!CW_w) * (init_ori_stackedprop+stckd_phase))

                # Adds the original polars that don't get copied in deepcopy
                this_prop.airfoils = copy_prop.airfoils
                this_prop._polars = copy_prop._polars
                this_prop._polarroot = copy_prop._polarroot
                this_prop._polartip = copy_prop._polartip

                push!(props_w_stacked_low, this_prop)
            end
        end
    end

    # Assembles fixed section of the wing
    main_wing_fixed = vlm.WingSystem()
    vlm.addwing(main_wing_fixed, "WingM", wing_md)
    if add_rotors
        for (i, prop) in enumerate(props_w_stacked_up)
            vlm.addwing(main_wing_fixed, "StackedPropUp$i", prop)
        end
        for (i, prop) in enumerate(props_w_stacked_low)
            vlm.addwing(main_wing_fixed, "StackedPropLow$i", prop)
        end
    end

    main_wing_R = vlm.WingSystem()
    vlm.addwing(main_wing_R, "Tip", wing_R)
    vlm.addwing(main_wing_R, "Winglet", winglet_R)

    main_wing_L = vlm.WingSystem()
    vlm.addwing(main_wing_L, "Tip", wing_L)
    vlm.addwing(main_wing_L, "Winglet", winglet_L)

    # Assembles moving sections of the wing
    main_wing_moving = vlm.WingSystem()
    vlm.addwing(main_wing_moving, "WingR", main_wing_R)
    vlm.addwing(main_wing_moving, "WingL", main_wing_L)
    if add_rotors
        for (i, prop) in enumerate(props_w)
            vlm.addwing(main_wing_moving, "Prop$i", prop)
        end
    end

    # offset to align with pivot line
    x_off_w = pivot_w*b_w/AR_w

    # Aligns moving and fixed section with their pivot line
    O_off_w = [-x_off_w, 0.0, 0.0]

    vlm.setcoordsystem(wing_md, O_off_w, Array(1.0I, 3, 3))

    for vlmwing in main_wing_moving.wings
        if typeof(vlmwing)==vlm.Rotor
            vlm.setcoordsystem(vlmwing, vlmwing._wingsystem.O + O_off_w,
                                        vlmwing._wingsystem.Oaxis; user=false)
        else
            vlm.setcoordsystem(vlmwing, vlmwing.O + O_off_w, vlmwing.Oaxis)
        end
    end

    # Places tilting sections at the main wing tip
    O_mv = (md_w*b_w/2)*[tan(lambda_w*pi/180), 0, tan(gamma_w*pi/180)]

    # Initial rotation of moving sections
    Oaxis_wmv = gt.rotation_matrix(0.0, -init_ori, 0.0)
    vlm.setcoordsystem(main_wing_moving, O_mv, Oaxis_wmv)

    # Assembles main wing
    main_wing = vlm.WingSystem()
    vlm.addwing(main_wing, "Fixed", main_wing_fixed)
    vlm.addwing(main_wing, "Moving", main_wing_moving)

    # Position of main wing
    O_w = [l_pos_w + x_off_w, 0, h_pos_w]
    Oaxis_w = gt.rotation_matrix(0.0, 0.0, 0.0)
    vlm.setcoordsystem(main_wing, O_w, Oaxis_w)

    # ------------ TANDEM WING -------------------------------------------
    if verbose; println("\t"^(v_lvl+1)*"Generating tandem wing assembly..."); end;

    # Generates tandem wing
    twing = vlm.simpleWing(b_tw, AR_tw, tr_tw, twist_r_tw, lambda_tw,
                                gamma_tw; twist_tip=twist_t_tw, n=n_tw, r=r_tw)

    # Middle section
    pos_md_tw = [-md_tw, 0.0, md_tw]
    clen_md_tw = [md_tw + (1/tr_tw)*(1-md_tw), 1/tr_tw,
                                                   md_tw + (1/tr_tw)*(1-md_tw)]
    twist_md_tw = [twist_r_tw + md_tw*(twist_t_tw-twist_r_tw),
                        twist_r_tw, twist_r_tw + md_tw*(twist_t_tw-twist_r_tw)]
    twing_md = vlm.complexWing(b_tw, AR_tw, ceil(Int, md_tw*n_tw), pos_md_tw,
                                  clen_md_tw, twist_md_tw,
                                  lambda_tw*ones(2), gamma_tw*ones(2);
                                  symmetric=false, chordalign=0.0,
                                  _ign1=true)
    # Left section
    pos_l_tw = [-1, pos_md_tw[1]]
    clen_l_tw = [1, clen_md_tw[1]]
    twist_l_tw = [twist_t_tw, twist_md_tw[1]]
    twing_L = vlm.complexWing(b_tw, AR_tw, ceil(Int, (1-md_tw)*n_tw/2), pos_l_tw,
                                    clen_l_tw, twist_l_tw,
                                    lambda_tw*ones(1), gamma_tw*ones(1);
                                    symmetric=false, chordalign=0.0,
                                    _ign1=true)
    # Right section
    pos_r_tw = [pos_md_tw[end], 1]
    clen_r_tw = [clen_md_tw[end], 1]
    twist_r_tw = [twist_md_tw[end], twist_t_tw]
    twing_R = vlm.complexWing(b_tw, AR_tw, ceil(Int, (1-md_tw)*n_tw/2), pos_r_tw,
                                    clen_r_tw, twist_r_tw,
                                    lambda_tw*ones(1), gamma_tw*ones(1);
                                    symmetric=false, chordalign=0.0,
                                    _ign1=true)

    ## Generates propellers on tandem wing (from right to left)
    if add_rotors
        if verbose; println("\t"^(v_lvl+2)*"Generating tandem wing propellers..."); end;
        O_prop_tw = [ ypos*[tan(lambda_tw*pi/180), 1, tan(gamma_tw*pi/180)] +
                        [-xoc_offset_tandem*AR_tw/b_tw, 0, 0]
                                                  for ypos in y_pos_prop_tw]
        props_tw = vlm.Rotor[]
        for i in 1:2*np_tw
            right = i<=np_tw    # Indicates which side of the wing
            copy_prop = tiltrotors[1+(i+(CW_tw!=CW_w))%2]
            this_prop = deepcopy(copy_prop)       # Alternates rotation orientation
            this_O = O_prop_tw[ right ? i : np_tw-(i-np_tw-1)] # Chooses position
            this_O = [1 0 0; 0 (-1)^!right 0; 0 0 1]*this_O   # Places it in correct side
            this_Oaxis = gt.rotation_matrix(0, (-1)^(i%2==0)*(-1)^right*tandem_pitchtilt, 0)

            vlm.setcoordsystem(this_prop, this_O, this_Oaxis; user=true)

            # Rotates props to be tip to tip
            vlm.rotate(this_prop, (-1)^(!CW_tw) * init_ori_prop)

            # Adds the original polars that don't get copied in deepcopy
            this_prop.airfoils = copy_prop.airfoils
            this_prop._polars = copy_prop._polars
            this_prop._polarroot = copy_prop._polarroot
            this_prop._polartip = copy_prop._polartip

            push!(props_tw, this_prop)
        end
    end

    # Assembles moving sections of the wing
    tandem_wing_moving = vlm.WingSystem()
    vlm.addwing(tandem_wing_moving, "WingR", twing_R)
    vlm.addwing(tandem_wing_moving, "WingL", twing_L)
    if add_rotors
        for (i, prop) in enumerate(props_tw)
            vlm.addwing(tandem_wing_moving, "Prop$i", prop)
        end
    end

    # offset to align with pivot line
    x_off_tw = pivot_tw*b_tw/AR_tw

    # Aligns moving and fixed section with their pivot line
    O_off_tw = [-x_off_tw, 0.0, 0.0]

    vlm.setcoordsystem(twing_md, O_off_tw, Array(1.0I, 3, 3))

    for vlmwing in tandem_wing_moving.wings
        if typeof(vlmwing)==vlm.Rotor
            vlm.setcoordsystem(vlmwing, vlmwing._wingsystem.O + O_off_tw,
                                        vlmwing._wingsystem.Oaxis; user=false)
        else
            vlm.setcoordsystem(vlmwing, vlmwing.O + O_off_tw, vlmwing.Oaxis)
        end
    end

    # Initial rotation of moving sections
    Oaxis_twmv = gt.rotation_matrix(0.0, -init_ori, 0.0)
    vlm.setcoordsystem(tandem_wing_moving, zeros(3), Oaxis_twmv)

    # Assembles tandem wing
    tandem_wing = vlm.WingSystem()
    vlm.addwing(tandem_wing, "FixedWing", twing_md)
    vlm.addwing(tandem_wing, "Moving", tandem_wing_moving)

    # Position of tandem wing
    O_tw = [l_pos_tw + x_off_tw, 0, h_pos_tw]
    Oaxis_tw = gt.rotation_matrix(0.0, 0.0, 0.0)
    vlm.setcoordsystem(tandem_wing, O_tw, Oaxis_tw)

    # ------------ FUSELAGE ----------------------------------------------
    if verbose; println("\t"^(v_lvl+1)*"Generating fuselage..."); end;
    fuselage = generatefuselage_vahana(l_f; ncells=20)

    body = gt.MultiGrid(3)
    gt.addgrid(body, "Fuselage", fuselage)

    if add_rotors
        if verbose; println("\t"^(v_lvl+1)*"Generating pylons..."); end;

        pylon_pos = 0.70
        pylon_length = stckd_xoc_offset*AR_w/b_w * pylon_pos

        for i in stacked
            pylon = generatepylon(pylon_length; ncells=5)

            right = i<=np_w
            this_O = O_prop_w[ right ? i : np_w-(i-np_w-1)]   # Chooses position
            this_O = [1 0 0; 0 (-1)^!right 0; 0 0 1]*this_O   # Places it in correct side
            this_O += main_wing.O                             # Translates it with the wing
                                                              # Offsets it according to length
            this_O += [1/3*0.5*pylon_pos, 0, 3*R_w*stckd_zoR_offset]
            rotation = gt.rotation_matrix2(0, 0, 0)

            gt.lintransform!(pylon, rotation, this_O)

            gt.addgrid(body, "Pylon$i", pylon)
        end
    end

    # ------------ SYSTEM ------------------------------------------------
    # Creates system assembly
    system = vlm.WingSystem()
    vlm.addwing(system, "MainWing", main_wing)
    vlm.addwing(system, "TandemWing", tandem_wing)

    if add_rotors
        rotors = vcat(props_w, props_w_stacked_up, props_w_stacked_low, props_tw)
    end

    # ------------ GROUND SURFACE ----------------------------------------
    ground1 = generate_ground_vahana(l_f)          # Take-off pad
    ground2 = generate_ground_vahana(l_f)          # Landing pad

    grounds = (ground1, ground2)

    # ------------ OUTPUTS------------------------------------------------
    # Tilting systems
    tilting_systems = (main_wing_moving, tandem_wing_moving)

    # Rotors grouped by systems of the same RPM
    if add_rotors
        rotor_systems = (props_w, props_w_stacked_up, props_w_stacked_low, props_tw)
    else
        rotor_systems = ()
    end

    # System to solve through the VLM solver
    vlm_system_m = vlm.WingSystem()
    vlm.addwing(vlm_system_m, "middle", wing_md)
    vlm.addwing(vlm_system_m, "R", wing_R)
    vlm.addwing(vlm_system_m, "letR", winglet_R)
    vlm.addwing(vlm_system_m, "L", wing_L)
    vlm.addwing(vlm_system_m, "letL", winglet_L)
    vlm_system_t = vlm.WingSystem()
    vlm.addwing(vlm_system_t, "R", twing_R)
    vlm.addwing(vlm_system_t, "L", twing_L)

    vlm_system = vlm.WingSystem()
    vlm.addwing(vlm_system, "MWing", vlm_system_m)
    vlm.addwing(vlm_system, "TWing", vlm_system_t)

    # Wake-shedding system (`vlm_system`+`rotors`)
    wake_system = vlm.WingSystem()
    vlm.addwing(wake_system, "SolveVLM", vlm_system)
    if add_rotors
        if VehicleType==uns.VLMVehicle
            for (i, rotor) in enumerate(rotors)
                vlm.addwing(wake_system, "Rotor$i", rotor)
            end
        else
            # Mute colinear warnings. This is needed since the quasi-steady solver
            #   will probe induced velocities at the lifting line of the blade
            uns.vlm.VLMSolver._mute_warning(true)
        end
    end

    # Dummy grids that are rotated and translated along with the vehicle
    grids = [body]

    # FVS's Vehicle object
    vehicle = VehicleType(   system;
                                tilting_systems=tilting_systems,
                                rotor_systems=rotor_systems,
                                vlm_system=vlm_system,
                                # wake_system=vlm_system,
                                wake_system=wake_system,
                                grids=grids
                             )

    return vehicle, grounds
end

"""
    Returns a loft of the Vahana fuselage.
"""
function generatefuselage_vahana(fuselage_length; ncells=20)
    sections = [[] for i in 1:6]

    # --------------- SECTION #1------------------------------------------------
    Ps = [
            [0.5, 0.1+0.2],
            [0.0, 1.0],
            [-0.5, 0.1+0.2],
            [0, -0.20+0.2]
         ]

    CPs = [
            [0.4, 1.0],
            [-0.4, 1.0],
            [-0.1, -0.3+0.2],
            [0.1, -0.3+0.2],
         ]

    rhos = [
            0.6,
            0.6,
            0.2,
            0.2
         ]

    scaling = 1.0

    ss = [
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells))
         ]

    points = gt.conic_cross_section(Ps, CPs, rhos, ss)
    points = scaling*points

    sections[1] = points

    # --------------- SECTION #2------------------------------------------------
    Ps = [
            [0.4, 0.4],
            [0.0, 1.0],
            [-0.4, 0.4],
            [0, -0.05]
         ]

    CPs = [
            [0.3, 1.0],
            [-0.3, 1.0],
            [-0.4, -0.08],
            [0.4, -0.08],
         ]

    rhos = [
            0.7,
            0.7,
            0.5,
            0.5
         ]

    scaling = 2.0

    ss = [
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells))
         ]

    points = gt.conic_cross_section(Ps, CPs, rhos, ss)
    points = scaling*points

    sections[2] = points;

    # --------------- SECTION #3------------------------------------------------
    Ps = [
            [0.8, 1.0],
            [0.0, 2.5],
            [-0.8, 1.0],
            [0, 0.2]
         ]

    CPs = [
            [0.6, 2.5],
            [-0.6, 2.5],
            [-0.8, 0.2],
            [0.8, 0.2],
         ]

    rhos = [
            0.7,
            0.7,
            0.5,
            0.5
         ]

    scaling = 1.0

    ss = [
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells))
         ]

    points = gt.conic_cross_section(Ps, CPs, rhos, ss)
    points = scaling*points

    sections[3] = points;

    # --------------- SECTION #4------------------------------------------------
    Ps = [
            [0.7, 2.0+0.25],
            [0.0, 3.0+0.25],
            [-0.7, 2.0+0.25],
            [0, 1.2]
         ]

    CPs = [
            [0.5, 3.0+0.25],
            [-0.5, 3.0+0.25],
            [-0.8, 1.2],
            [0.8, 1.2],
         ]

    rhos = [
            0.7,
            0.7,
            0.5,
            0.5
         ]

    scaling = 1.0

    ss = [
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells))
         ]

    points = gt.conic_cross_section(Ps, CPs, rhos, ss)
    points = scaling*points

    sections[4] = points;

    # --------------- SECTION #5------------------------------------------------
    Ps = [
            [0.4, 2.75-0.2],
            [0.0, 3.25-0.2],
            [-0.4, 2.75-0.2],
            [0, 1.7]
         ]

    CPs = [
            [0.3, 3.25-0.2],
            [-0.3, 3.25-0.2],
            [-0.4, 1.7],
            [0.4, 1.7],
         ]

    rhos = [
            0.4,
            0.4,
            0.5,
            0.5
         ]

    scaling = 1.0

    ss = [
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells))
         ]

    points = gt.conic_cross_section(Ps, CPs, rhos, ss)
    points = scaling*points

    sections[5] = points;

    # --------------- SECTION #6------------------------------------------------
    Ps = [
            [0.4-0.2, 3.25],
            [0.0, 3.5-0.1],
            [-0.4+0.2, 3.25],
            [0, 2.5+0.5]
         ]

    CPs = [
            [0.3-0.2, 3.5-0.1],
            [-0.3+0.2, 3.5-0.1],
            [-0.4, 2.5+0.5],
            [0.4, 2.5+0.5],
         ]

    rhos = [
            0.4,
            0.4,
            0.3,
            0.3
         ]

    scaling = 1.0

    ss = [
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells))
         ]
    points = gt.conic_cross_section(Ps, CPs, rhos, ss)
    points = scaling*points

    sections[6] = points;

    # --------------- SECTION #0------------------------------------------------
    Ps = [
            [0.5, 0.1+0.2],
            [0.0, 1.0],
            [-0.5, 0.1+0.2],
            [0, -0.20+0.2]
         ]

    CPs = [
            [0.4, 1.0],
            [-0.4, 1.0],
            [-0.1, -0.3+0.2],
            [0.1, -0.3+0.2],
         ]

    rhos = [
            0.6,
            0.6,
            0.2,
            0.2
         ]

    scaling = 0.1

    ss = [
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells))
         ]
    points = gt.conic_cross_section(Ps, CPs, rhos, ss)
    points = scaling*points

    section0 = points

    # --------------- LOFT -----------------------------------------------------
    # fuselage_length = 5.86              # (m) fuselage length
    bscale = fuselage_length            # Scales the entire geometry by this factor

    crosssections = [ # (l/bscale position , contour)
                        (0.0  , [[2.00*0.8, 0.25][k]*(point[k]) + [0, 0.3][k] for point in section0, k in 1:2]),
                        (0.02, [[0.6*0.8, 0.3][k]*(point[k]) + [0, 0.15][k] for point in sections[1], k in 1:2]),
                        (1/6*0.5, [[1.0*0.8, 0.75][k]*point[k] for point in sections[1], k in 1:2]),
    #                     (1/6*0.5, [[0.75, 0.75][k]*point[k] for point in sections[1], k in 1:2]),
                        (2/6, [[1.00*0.8, 0.75][k]*point[k] for point in sections[2], k in 1:2]),
                        (3/6, [[1.00*0.8, 0.75][k]*(point[k] - [0, 0.1*2][k]) for point in sections[3], k in 1:2]),
                        (4/6, [[0.75, 0.75][k]*(point[k] - [0, 0.4*2][k]) for point in sections[4], k in 1:2]),
                        (5/6, [[0.75, 0.75][k]*(point[k] - [0, 0.6][k]) for point in sections[5], k in 1:2]),
                        (6/6, [[0.75, 0.75][k]*(point[k] - [0, 1.0][k]) for point in sections[6], k in 1:2]),
                    ]

    # Dummy parameters
    b_pos = [sec[1] for sec in crosssections]
    chords = (1/bscale)*ones(size(b_pos))
    twists = zeros(size(b_pos))
    LE_x = zeros(size(b_pos))
    LE_z = zeros(size(b_pos))

    tilt_y = zeros(size(b_pos))
    # tilt_y[2] = 45

    grid_fuselage = gt.generate_loft(crosssections, bscale,
                                        b_pos, chords, twists, LE_x, LE_z;
                                        # MORE GEOMETRIC OPTIONS
                                        tilt_y=tilt_y,
                                        symmetric=false,
                                        loop_dim=0,
                                    )

    # Aligns the grid with the x-axis
    rotation = gt.rotation_matrix2(0, 0, 90)
    gt.lintransform!(grid_fuselage, rotation, zeros(3))

    dimsplit = 1
    triangrid_fuselage = gt.GridTriangleSurface(grid_fuselage, dimsplit)

    return triangrid_fuselage
end

"""
    Returns a loft of the pylon.
"""
function generatepylon(pylon_length; ncells=20)
    sections = [[] for i in 1:3]

    # --------------- SECTION #1------------------------------------------------
    Ps = [
            [0.5, 0.1+0.2],
            [0.0, 1.0],
            [-0.5, 0.1+0.2],
            [0, -0.20+0.2]
         ]

    CPs = [
            [0.4, 1.0],
            [-0.4, 1.0],
            [-0.1, -0.3+0.2],
            [0.1, -0.3+0.2],
         ]

    rhos = [
            0.6,
            0.6,
            0.2,
            0.2
         ]

    scaling = 1.0

    ss = [
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells))
         ]

    points = gt.conic_cross_section(Ps, CPs, rhos, ss)
    points = scaling*points

    sections[1] = points

    # --------------- SECTION #2------------------------------------------------
    Ps = [
            [0.4, 0.4],
            [0.0, 1.0],
            [-0.4, 0.4],
            [0, -0.05]
         ]

    CPs = [
            [0.3, 1.0],
            [-0.3, 1.0],
            [-0.4, -0.08],
            [0.4, -0.08],
         ]

    rhos = [
            0.7,
            0.7,
            0.5,
            0.5
         ]

    scaling = 2.0

    ss = [
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells))
         ]

    points = gt.conic_cross_section(Ps, CPs, rhos, ss)
    points = scaling*points

    sections[2] = points;

    # --------------- SECTION #3------------------------------------------------
    Ps = [
            [0.4, 2.75-0.2],
            [0.0, 3.25-0.2],
            [-0.4, 2.75-0.2],
            [0, 1.7]
         ]

    CPs = [
            [0.3, 3.25-0.2],
            [-0.3, 3.25-0.2],
            [-0.4, 1.7],
            [0.4, 1.7],
         ]

    rhos = [
            0.4,
            0.4,
            0.5,
            0.5
         ]

    scaling = 1.0

    ss = [
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells))
         ]

    points = gt.conic_cross_section(Ps, CPs, rhos, ss)
    points = scaling*points

    points = [p .- [0, 1.7] for p in points]
    # points = [p .* [1, 0.75] for p in points]

    sections[3] = points;

    # --------------- SECTION #0------------------------------------------------
    Ps = [
            [0.5, 0.1+0.2],
            [0.0, 1.0],
            [-0.5, 0.1+0.2],
            [0, -0.20+0.2]
         ]

    CPs = [
            [0.4, 1.0],
            [-0.4, 1.0],
            [-0.1, -0.3+0.2],
            [0.1, -0.3+0.2],
         ]

    rhos = [
            0.6,
            0.6,
            0.2,
            0.2
         ]

    scaling = 0.1

    ss = [
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells)),
            collect(range(0, 1, length=ncells))
         ]
    points = gt.conic_cross_section(Ps, CPs, rhos, ss)
    points = scaling*points

    section0 = points

    # --------------- LOFT -----------------------------------------------------

    # --------------- LOFT -----------------------------------------------------
    # fuselage_length = 5.86              # (m) fuselage length
    bscale = pylon_length            # Scales the entire geometry by this factor

    crosssections = [ # (l/bscale position , contour)
                    (0.0  , [[2.00*0.8, 0.25][k]*(point[k]) + [0, 0.3][k] for point in section0, k in 1:2]*pylon_length/5.86),
                    (0.02, [[0.6*0.8, 0.3][k]*(point[k]) + [0, 0.15][k] for point in sections[1], k in 1:2]*pylon_length/5.86),
                    (1/3*0.25, [[1.0*0.8, 0.75*0.75][k]*point[k] for point in sections[1], k in 1:2]*pylon_length/5.86),
                    (2/3, [[1.00*0.8, 0.35*0.75][k]*point[k] for point in sections[2], k in 1:2]*pylon_length/5.86),
                    (3/3, [[1.00*0.8, 0.75*0.75][k]*(point[k] - [0, 0.1*2][k]) for point in sections[3], k in 1:2]*pylon_length/5.86),
                    ]

    # Dummy parameters
    b_pos = [sec[1] for sec in crosssections]
    chords = (1/bscale)*ones(size(b_pos))
    twists = zeros(size(b_pos))
    LE_x = zeros(size(b_pos))
    LE_z = zeros(size(b_pos))

    tilt_y = zeros(size(b_pos))
    # tilt_y[2] = 45

    grid_pylon = gt.generate_loft(crosssections, bscale,
                                        b_pos, chords, twists, LE_x, LE_z;
                                        # MORE GEOMETRIC OPTIONS
                                        tilt_y=tilt_y,
                                        symmetric=false,
                                        loop_dim=0,
                                    )

    # Aligns the grid with the x-axis
    rotation = gt.rotation_matrix2(0, 0, 90)
    gt.lintransform!(grid_pylon, rotation, zeros(3))

    dimsplit = 1
    triangrid_pylon = gt.GridTriangleSurface(grid_pylon, dimsplit)

    return triangrid_pylon
end

"""
    Generates ground surface.
"""
function generate_ground_vahana(bscale; verify_spline=false, spl_s=0.001,
                                spl_k="automatic", visualize=false)

    ## Ellipse
    fx, fy = 1.7, 1.05
    lx = bscale*3.5
    ly = bscale*3
    cx = -lx*3/4
    cy = 0.0
    cz = -0.06*bscale
    a, b = lx, ly
    perimeter = [ a*b/sqrt((b*cos(tht))^2 + (a*sin(tht))^2
                    ) * [fx*cos(tht), fy*sin(tht), 0] + [cx, cy, cz] for
                                            tht in range(0, 2*pi, length=179)]

    # Fluid domain
    NDIVSx = 50              # Cells in the parametric x-direction
    NDIVSy = 50              # Cells in the parametric y-direction
    NDIVSz = 0               # Cells in the geometric z-direction

    z_off = cz               # z-offset for ground position

    if visualize
        opt_args = [(:save_path, "temps/"), (:file_name, "vahanaground"),
                    (:paraview, true)]
    else
        opt_args = []
    end

    ground = generate_perimetergrid(perimeter, NDIVSx, NDIVSy, NDIVSz;
                                       z_min=z_off, z_max=z_off,
                                        verify_spline=verify_spline, spl_s=spl_s,
                                        spl_k=spl_k,
                                        opt_args...
                                        )

    return ground
end

# -------------- END OF VAHANA -------------------------------------------------

################################################################################
# UTILITIES
################################################################################

"""
`generate_perimetergrid(perimeter::Array{Array{T, 1}, 1},
                                  NDIVSx, NDIVSy, NDIVSz;
                                  z_min::Real=0, z_max::Real=0,
                                  # SPLINE OPTIONS
                                  verify_spline::Bool=true,
                                  spl_s=0.001, spl_k="automatic",
                                  # FILE OPTIONS
                                  save_path=nothing, file_name="perimeter",
                                  paraview=true
                                )`

  Generates the perimeter grid with `perimeter` the array of points of the
  contour (must be a closed contour), and `NDIVS_` the number of cells in
  each parametric dimension (give it `NDIVSz=0` for a flat surface, otherwise
  it'll generate a volumetric grid between `z_min` and `z_max`).
"""
function generate_perimetergrid(perimeter::Array{Array{T, 1}, 1},
                                  NDIVSx, NDIVSy, NDIVSz;
                                  z_min::Real=0, z_max::Real=0,
                                  # SPLINE OPTIONS
                                  verify_spline::Bool=true,
                                  spl_s=0.001, spl_k="automatic",
                                  # FILE OPTIONS
                                  save_path=nothing, file_name="perimeter",
                                  paraview=true
                                ) where{T<:Real}

  # Error cases
  multidiscrtype = Array{Tuple{Float64,Int64,Float64,Bool},1}
  if typeof(NDIVSx)==Int64
    nz = NDIVSz
  elseif typeof(NDIVSz)==multidiscrtype
    nz = 0
    for sec in NDIVSz
      nz += sec[2]
    end
  else
    error("Expected `NDIVSz` to be type $(Int64) or $MultiDiscrType,"*
            " got $(typeof(NDIVSz)).")
  end

  # --------- REPARAMETERIZES THE PERIMETER ---------------------------
  org_x = [p[1] for p in perimeter]
  org_y = [p[2] for p in perimeter]
  # Separate upper and lower sides to make the contour injective in x
  upper, lower = gt.splitcontour(org_x, org_y)

  # # Parameterize both sides independently
  # fun_upper = gt.parameterize(upper[1], upper[2], zeros(upper[1]); inj_var=1,
  #                                                     s=spl_s, kspl=spl_k)
  # fun_lower = gt.parameterize(lower[1], lower[2], zeros(lower[1]); inj_var=1,
  #                                                     s=spl_s, kspl=spl_k)
  # # Discretizes both sides
  # if NDIVSx==multidiscrtype
  #   new_upper = gt.multidiscretize(fun_upper, 0, 1, NDIVSx)
  #   new_lower = gt.multidiscretize(fun_lower, 0, 1, NDIVSx)
  # else
  #   new_upper = gt.discretize(fun_upper, 0, 1, NDIVSx, 1.0)
  #   new_lower = gt.discretize(fun_lower, 0, 1, NDIVSx, 1.0)
  # end

  splt_up = Int(ceil((size(upper[1],1)/2)))
  splt_low = Int(ceil((size(lower[1],1)/2)))

  # Splits the perimeter into four faces
  upper1 = [[x for x in upper[1][1:splt_up]], [y for y in upper[2][1:splt_up]]]
  upper2 = [[x for x in upper[1][splt_up:end]], [y for y in upper[2][splt_up:end]]]
  lower1 = [[x for x in lower[1][1:splt_low]], [y for y in lower[2][1:splt_low]]]
  lower2 = [[x for x in lower[1][splt_low:end]], [y for y in lower[2][splt_low:end]]]

  # Parameterize both sides independently
  fun_upper1 = gt.parameterize(upper1[1], upper1[2], zeros(size(upper1[1])); inj_var=1,
                                                      s=spl_s, kspl=spl_k)
  fun_upper2 = gt.parameterize(upper2[1], upper2[2], zeros(size(upper2[1])); inj_var=1,
                                                      s=spl_s, kspl=spl_k)
  fun_lower1 = gt.parameterize(lower1[1], lower1[2], zeros(size(lower1[1])); inj_var=1,
                                                      s=spl_s, kspl=spl_k)
  fun_lower2 = gt.parameterize(lower2[1], lower2[2], zeros(size(lower2[1])); inj_var=1,
                                                      s=spl_s, kspl=spl_k)
  # Discretizes both sides
  if NDIVSx==multidiscrtype
    new_upper1 = gt.multidiscretize(fun_upper1, 0, 1, NDIVSx)
    new_upper2 = gt.multidiscretize(fun_upper2, 0, 1, NDIVSy)
    new_lower1 = gt.multidiscretize(fun_lower1, 0, 1, NDIVSy)
    new_lower2 = gt.multidiscretize(fun_lower2, 0, 1, NDIVSx)
  else
    new_upper1 = gt.discretize(fun_upper1, 0, 1, NDIVSx, 1.0)
    new_upper2 = gt.discretize(fun_upper2, 0, 1, NDIVSy, 1.0)
    new_lower1 = gt.discretize(fun_lower1, 0, 1, NDIVSy, 1.0)
    new_lower2 = gt.discretize(fun_lower2, 0, 1, NDIVSx, 1.0)
  end

  # ----------------- SPLINE VERIFICATION --------------------------------------

  if verify_spline
  #   new_points = vcat(reverse(new_upper), new_lower)
    # new_x = [p[1] for p in new_points]
    # new_y = [p[2] for p in new_points]
    plt.plot(org_x, org_y, "--^k", label="Original", alpha=0.5)
    # plt.plot(new_x, new_y, ":.b", label="Parameterized")

    plt.plot(upper1[1], upper1[2], "*b", label="upper1", alpha=0.75)
    plt.plot(upper2[1], upper2[2], "^b", label="upper2", alpha=0.75)
    plt.plot(lower1[1], lower1[2], "*r", label="lower1", alpha=0.75)
    plt.plot(lower2[1], lower2[2], "^r", label="lower2", alpha=0.75)

    new_points = vcat(new_upper1, new_upper2)
    new_x = [p[1] for p in new_points]
    new_y = [p[2] for p in new_points]
    plt.plot(new_x, new_y, ":.b", label="Parameterized Upper", alpha=0.75)
    new_points = vcat(new_lower1, new_lower2)
    new_x = [p[1] for p in new_points]
    new_y = [p[2] for p in new_points]
    plt.plot(new_x, new_y, ":.r", label="Parameterized Lower", alpha=0.75)
    # plt.xlabel(plt.L"x")
    # plt.ylabel(plt.L"y")
    plt.legend(loc="best")
    plt.grid(true, color="0.8", linestyle="--")
  end

  # --------- GRIDS THE INSIDE OF THE PERIMETER ---------------------
  # Parametric grid
  P_min = zeros(3)
  P_max = [1, 1, 1*(nz!=0)]
  param_grid = gt.Grid(P_min, P_max, [NDIVSx, NDIVSy, NDIVSz])

  # function my_space_transform(X, ind)
  #     i = ind[1]                      # Arc length point
  #     w = X[2]                        # Weight
  #     z = z_min + X[3]*(z_max-z_min)  # z-position
  #
  #     Y = new_lower[i] + w*(new_upper[i]-new_lower[i])
  #     Y[3] = z
  #
  #     return Y
  # end

  rev_new_upper1 = reverse(new_upper1) # Left edge
  rev_new_upper2 = reverse(new_upper2) # Upper edge
  rev_new_lower1 = reverse(new_lower1) # Lower edge
  rev_new_lower2 = reverse(new_lower2) # Right edge

  function my_space_transform(X, ind)
      beta = 1e-2*pi/2                 # Warping parameter

      z = z_min + X[3]*(z_max-z_min)   # z-position

      xw = X[2]                        # x weight
      yw = X[1]                        # y weight

      # Linearly interpolate the edges
      x = new_upper1[ind[1]]*xw + new_lower2[ind[1]]*(1-xw) # Left-right edges
      y = rev_new_upper2[ind[2]]*yw + rev_new_lower1[ind[2]]*(1-yw) # Upper-lower edges
      # Y = x + y

      # A clever way of merging the two interpolations
      # First weight
      yw = abs(X[1]-0.5)/0.5
      yw = tan(yw*pi/2 - beta)/tan(pi/2 - beta)
      # Second weight
      xw = abs(X[2]-0.5)/0.5
      xw = tan(xw*pi/2 - beta)/tan(pi/2 - beta)
      # Merged weight
      w = ( (1-xw) + yw ) /2
      # Merged point
      Y = x*(1-w) + y*w

      Y[3] = z

      return Y
  end

  # Applies the space transformation to the parametric grid
  gt.transform3!(param_grid, my_space_transform)

  if save_path!=nothing
    gt.save(param_grid, file_name; path=save_path)

    if paraview
      strn = file_name*".vtk"
      run(`paraview --data=$save_path/$strn`)
    end

  end

  return param_grid::gt.Grid
end
# -------------- END OF COMMONS ------------------------------------------------


using Wflow
using Dates
using TOML

# edit existing TOML to get a short run in a temporary directory
tomlpath = joinpath(@__DIR__, "sbm_config.toml")
confdict = TOML.parsefile(tomlpath)
confdict["endtime"] = DateTime("2000-01-03T00:00:00")
# convert relative paths to absolute paths, to avoid having to copy input files
confdict["input"]["path_forcing"] = joinpath(@__DIR__, confdict["input"]["path_forcing"])
confdict["input"]["path_static"] = joinpath(@__DIR__, confdict["input"]["path_static"])
tmp_dir = mktempdir()
tomlpath = joinpath(tmp_dir, "config.toml")
open(tomlpath, "w") do io
    TOML.print(io, confdict)
end

Wflow.run_simulation(tomlpath)

#=
# test whether restarted runs get the same results as continuus ones, i.e. state is captured
tomlpath = joinpath(@__DIR__, "sbm_config.toml")
config = Wflow.Config(tomlpath)

# January at once, cold start
config.starttime = DateTime("2000-01-01T00:00:00")
config.endtime   = DateTime("2000-02-01T00:00:00")
config.model.reinit = true  # cold start
# note that this needs to be relative to the tomlpath
config.state.path_output = joinpath(dirname(tomlpath), "data/state-test/outstates-moselle-january.nc")
config.output.path = joinpath(dirname(tomlpath), "data/state-test/output-moselle-january.nc")
model = Wflow.initialize_sbm_model(config)
Wflow.run_simulation(model)

# first half of January, cold start
config.starttime = DateTime("2000-01-01T00:00:00")
config.endtime   = DateTime("2000-01-15T00:00:00")
config.model.reinit = true  # cold start
config.state.path_output = joinpath(dirname(tomlpath), "data/state-test/outstates-moselle-january-1of2.nc")
config.output.path = joinpath(dirname(tomlpath), "data/state-test/output-moselle-january-1of2.nc")
model = Wflow.initialize_sbm_model(config)
Wflow.run_simulation(model)

# second half of January, warm start
config.starttime = DateTime("2000-01-15T00:00:00")
config.endtime   = DateTime("2000-02-01T00:00:00")
config.model.reinit = false  # cold start
config.state.path_input = joinpath(dirname(tomlpath), "data/state-test/outstates-moselle-january-1of2.nc")
config.state.path_output = joinpath(dirname(tomlpath), "data/state-test/outstates-moselle-january-2of2.nc")
config.output.path = joinpath(dirname(tomlpath), "data/state-test/output-moselle-january-2of2.nc")
model = Wflow.initialize_sbm_model(config)
Wflow.run_simulation(model)
=#


module Reckoner

    include("interface.jl")

    export AbstractMatch, AbstractMatches

    export DefaultMatch, DefaultMatches

    export ReckonerInstance

    export reckoner_defaults

    export  aup,
            weight,
            challenge_window,
            skill,
            rating,
            eff_challenge,
            win_chances,
            display_rank,
            player_win_chances,
            weights,
            ratings,
            challenge_windows,
            skills

    export  challenge,
            timestamp,
            win,
            team_id

end # module


using Base.Test

@testset "abstract base types" begin
	@test AbstractInterpolation() == AbstractInterpolation{1}
	@test AbstractInterpolation(3) == AbstractInterpolation{3}
	@test ndims(AbstractInterpolation{3}) == 3

	@test CompoundInterpolation{3, Tuple{AbstractInterpolation{3}, AbstractInterpolation{4}}} <: AbstractInterpolation{3}
	@test CompoundInterpolation(Tuple{AbstractInterpolation{2}, AbstractInterpolation{4}}) == CompoundInterpolation{6, Tuple{AbstractInterpolation{2}, AbstractInterpolation{4}}}
	@test CompoundInterpolation(AbstractInterpolation(), AbstractInterpolation{4}) == CompoundInterpolation{5, Tuple{AbstractInterpolation{1}, AbstractInterpolation{4}}}
	@test ndims(CompoundInterpolation{6, Tuple{AbstractInterpolation{2}, AbstractInterpolation{4}}}) == 6
end;


struct MultiFieldCellField{R,S} <: CellField
  single_fields::Vector{<:CellField}
  cell_map::AbstractArray
  ref_trait::Val{R}
  cell_axes::AbstractArray
  metasize::Val{S}
  function MultiFieldCellField(single_fields::Vector{<:CellField})
    @assert length(single_fields) > 0
    f1 = first(single_fields)
    cell_map = get_cell_map(f1)
    ref_style = RefStyle(f1)
    cell_axes = get_cell_axes(f1)
    metasize = MetaSizeStyle(f1)
    R = get_val_parameter(ref_style)
    S = get_val_parameter(metasize)
    new{R,S}(single_fields,cell_map,ref_style,cell_axes,metasize)
  end
end

Arrays.get_array(a::MultiFieldCellField) = @notimplemented

CellData.get_memo(a::MultiFieldCellField) = @notimplemented

CellData.get_cell_map(a::MultiFieldCellField) = a.cell_map

CellData.get_cell_axes(a::MultiFieldCellField) = a.cell_axes

CellData.RefStyle(::Type{<:MultiFieldCellField{R}}) where R = Val(R)

CellData.MetaSizeStyle(::Type{<:MultiFieldCellField{R,S}}) where {R,S} = Val(S)

Base.length(f::MultiFieldCellField) = length(first(f.single_fields))

num_fields(a::MultiFieldCellField) = length(a.single_fields)

Base.getindex(a::MultiFieldCellField,i) = a.single_fields[i]

Base.iterate(a::MultiFieldCellField)  = iterate(a.single_fields)

Base.iterate(a::MultiFieldCellField,state)  = iterate(a.single_fields,state)

function Fields.evaluate(cf::MultiFieldCellField,x::AbstractArray)
  s = """
  Evaluation of MultiFieldCellField objects is not implemented at this moment.
  You need to extract the individual fields and then evaluate them separatelly.
  If ever implemented, evaluating a `MultiFieldCellField` directly would provide,
  at each evaluation point, a tuple with the value of the different fields.
  """
  @notimplemented s
end

function Geometry.restrict(a::MultiFieldCellField,trian::Triangulation)
  f = (ai) -> restrict(ai,trian)
  blocks = map(f,a.single_fields)
  blocks
end


using DiffEqFlux, Flux

"""
    Module that modifies parameters
"""
abstract type Modifier <: DiffEqFlux.FastLayer end

""" Required as a FastLayer, define a paramsum field """
paramlength(f::Modifier) = f.param_size

""" Required as a FastLayer """
initial_params(f::Modifier) = f.initial_params()

""" Nominal modifier does not change the parameters """
struct Nominal{P} <: Modifier
    name::String
    param_size::Int
    initial_params::P
    function Nominal(init = Flux.zeros)
        initial_params() = init(0)
        new{typeof(initial_params)}("Nominal", 0, initial_params)
    end
end

function (f::Nominal)(p_sub, p_link, p)
    return p_sub, p_link
end

function debug(f::Nominal, p_sub, p_link, p)
    return 0.0
end

include("iss_ctrnn0.jl")
include("iss_ctrnn1.jl")


include("../imports.jl")

@testset ExtendedTestSet "`ComposedProjectiveTransform`" begin
    tfms = (Project(Translation(10, 10)), Project(Translation(-10, 10)))
    tfm = ComposedProjectiveTransform(tfms[1], tfms[2])
    item = Image(rand(10, 10))
    ## composition of `AbstractProjective`s creates a `ComposedProjectiveTransform`
    @test tfm == tfms[1] |> tfms[2]
    ## can be `apply`ed
    @test_nowarn apply(tfm, item)

    ## identical result to sequential application
    @test itemdata(apply(tfm, item)) == itemdata(apply(Sequence(tfms), item))
end


#     using MAT

    struct NCSSeg
        qwTimeStamp::UInt64
        dwChannelNumber::UInt32
        dwSampleFreq::UInt32
        dwNumValidSamples::UInt32
        snSample::Array{Int16,1}
    end

    struct NCSData
        filename::String
        header::String
        seg_num::UInt32
        data::Array{NCSSeg, 1}
        sample::Array{Int16, 1}
    end

    # read neuralynx file
    function readNCSFile(filename)

        f = open(filename)
        header = read(f, UInt8, 16*1024) .|> Char |> String
        bdata = reinterpret(Int16, read(f))
        close(f)

        records = length(bdata) / 522 |> Int
        data = Array{NCSSeg, 1}(records)

        for i = 1:records
            segdata = bdata[1+(i-1)*522:522*i]

            timestamp = reinterpret(UInt64, segdata[1:4])
            chnum = reinterpret(UInt32, segdata[5:6])
            samplefreq = reinterpret(UInt32, segdata[7:8])
            samplenum = reinterpret(UInt32, segdata[9:10])

            data[i] = NCSSeg(timestamp[1], chnum[1], samplefreq[1], samplenum[1], segdata[11:522])
        end

        sample = zeros(Int16, records*512)
        for idx = 1:records
            sample[1+512*(idx-1):512*idx] = data[idx].snSample
        end

        return NCSData(filename, header, records, data, sample)
    end

    # # save into mat-v7.3 file
    # function save(filename::String, ncs::NCSData; mode=:sampleonly)
    #
    #     if mode == :sampleonly
    #         matwrite(filename, Dict("header"=>ncs.header, "sample"=>ncs.sample))
    #     elseif mode == :all
    #         matwrite(filename,
    #             Dict("header"=>ncs.header,
    #                 "seg_num"=>ncs.seg_num,
    #                 "data"=>ncs.data,
    #                 "sample"=>ncs.sample))
    #     else
    #         matwrite(filename, Dict("header"=>ncs.header, "sample"=>ncs.sample))
    #     end
    # end


# convert_a3d

# Base definitions of method to convert draws to the desired `output_format`.

"""

Convert the output file created by cmdstan to the shape of choice

$(SIGNATURES)

### Method
```julia
convert_a3d(a3d_array, cnames; output_format=::Val{Symbol}, start=1)
```
### Required arguments
```julia
* `a3d_array::Array{Float64, 3},`      : Read in from output files created by cmdstan
* `cnames::Vector{AbstractString}`     : Monitored variable names
```

### Optional arguments
```julia
* `::Val{Symbol}`                      : Output format, default is :mcmcchains
* `::start=1`                          : First draw for MCMCChains.Chains
```

### Return values
```julia
* `res`                       : Draws converted to the specified format.
```

"""
convert_a3d(a3d_array, cnames, ::Val{:array}; kwargs...) = a3d_array


include("Chipmunk.jl")

space = Chipmunk.cpSpaceNew()
Chipmunk.cpSpaceSetGravity(space, Chipmunk.cpVect(0, 10))
body = Chipmunk.cpBodyNew(1, 1)
Chipmunk.cpSpaceAddBody(space, body)

for i = 1:10
    Chipmunk.cpSpaceStep(space, 1 / 60)

    p = Chipmunk.cpBodyGetPosition(body)
    v = Chipmunk.cpBodyGetVelocity(body)

    println("position = ($(p.x), $(p.y)), velocity = ($(v.x), $(v.y))")
end


__precompile__()
module Remove_response_obspy

using PyCall, SeisIO, Suppressor

export remove_response_obspy, remove_response_obspy!

"""
remove_response_obspy(S::SeisData, xf::String; pre_filt::NTuple{4,Float64}=(0.001, 0.002, 10.0, 20.0), zeropadlen::Float64 = 30*60, output::String="VEL")
remove Instrumental response with obspy.
Currently available only one channel.
This function will be replaced by julia original function

#Input
 - zeropadlen: [min] zeropad at the side of dataset to avoid edge effect

"""
function remove_response_obspy!(S::SeisData, xf::String; pre_filt::NTuple{4,Float64}=(0.001, 0.002, 10.0, 20.0), zeropadlen::Float64 = 30*60.0, output::String="VEL")

    @suppress begin
    Stream          = pyimport_conda("obspy.core.stream", "Stream")
    Trace           = pyimport_conda("obspy.core.trace", "Trace")
    read_inventory  = pyimport_conda("obspy", "read_inventory")
    UTCDateTime     = pyimport_conda("obspy.core.utcdatetime", "UTCDateTime")

    for i = 1:length(S.x)
        #add zero at the edge of data to avoid edge effect
        numofzeropad = round(Int64, zeropadlen * S.fs[i])
        #z1 = zeros(numofzeropad)
        #---modify 2019/06/11: add fliped signal to avoid edge effect
        #SeisIO.taper!(S)

        if length(S.x[i]) < numofzeropad
            numofzeropad = round(Int64, length(S.x[i])/2)
        end

        zl = S.x[i][numofzeropad:-1:1]
        zr = S.x[i][end:-1:end-numofzeropad+1]
        #z1 = zeros(numofzeropad)
        #x_withpad = vcat(z1,S.x[i],z1)
        x_withpad = vcat(zl,S.x[i],zr)
        trace=Trace.Trace()
        #trace.data = S.x[i]
        trace.data = x_withpad
        trace.stats.sampling_rate = S.fs[i]
        trace.stats.delta = 1.0./ S.fs[i]
        trace.stats.starttime = UTCDateTime.UTCDateTime(string(u2d(S.t[i][1,2]*1e-6)))
        str1 = split(S.id[i], ".")
        trace.stats.network = str1[1]
        trace.stats.station=  str1[2]
        trace.stats.location= str1[3]
        trace.stats.channel=  str1[4]

        stream = Stream.Stream(traces=trace)
        #read stationXML
        inv = read_inventory.read_inventory(xf)
        #stresm.remove_sensitivity(inventory=inv)
        stream.remove_response(inventory=inv, output=output, pre_filt=pre_filt, plot=false)
        x_withpad_removed =  stream.traces[1].data
        S.x[i] = x_withpad_removed[numofzeropad+1:end-numofzeropad]
        #add note
        SeisIO.note!(S[i], "remove_response_obspy!, pre_filt=$pre_filt")
    end

    end
    return nothing
end

function remove_response_obspy(S::SeisData, xf::String; pre_filt::NTuple{4,Float64}=(0.001, 0.002, 10.0, 20.0), zeropadlen::Float64 = 30*60.0, output::String="VEL")
  U = deepcopy(S)
  remove_response_obspy!(U, pre_filt=pre_filt,  zeropadlen, output=output)
  return U
end

end


##### deal with nutrients uptake
function gpu_calc_consume_kernel!(ctsdic, ctsdoc, ctsnh4, ctsno3, ctspo4, plank, ac, x, y, z, ΔT)
    index = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    stride = blockDim().x * gridDim().x
    for i = index:stride:size(ac,1)
        @inbounds CUDA.@atomic ctsdic[x[i], y[i], z[i]] += (plank.resp[i] - plank.PS[i]) * ΔT * ac[i]
        @inbounds CUDA.@atomic ctsdoc[x[i], y[i], z[i]] += (plank.exu[i] - plank.VDOC[i]) * ΔT * ac[i]
        @inbounds CUDA.@atomic ctsnh4[x[i], y[i], z[i]] += -plank.VNH4[i] * ΔT * ac[i]
        @inbounds CUDA.@atomic ctsno3[x[i], y[i], z[i]] += -plank.VNO3[i] * ΔT * ac[i]
        @inbounds CUDA.@atomic ctspo4[x[i], y[i], z[i]] += -plank.VPO4[i] * ΔT * ac[i]
    end
    return nothing
end
function calc_consume!(ctsdic, ctsdoc, ctsnh4, ctsno3, ctspo4, plank, ac, x, y, z, ΔT, ::GPU)
    @cuda threads=256 blocks=ceil(Int, size(ac,1)/256) gpu_calc_consume_kernel!(ctsdic, ctsdoc, ctsnh4,
                                                        ctsno3, ctspo4, plank, ac, x, y, z, ΔT)
    return nothing
end
function calc_consume!(ctsdic, ctsdoc, ctsnh4, ctsno3, ctspo4, plank, ac, x, y, z, ΔT, ::CPU)
    for i in 1:size(ac,1)
        @inbounds ctsdic[x[i], y[i], z[i]] += (plank.resp[i] - plank.PS[i]) * ΔT * ac[i]
        @inbounds ctsdoc[x[i], y[i], z[i]] += (plank.exu[i] - plank.VDOC[i]) * ΔT * ac[i]
        @inbounds ctsnh4[x[i], y[i], z[i]] += -plank.VNH4[i] * ΔT * ac[i]
        @inbounds ctsno3[x[i], y[i], z[i]] += -plank.VNO3[i] * ΔT * ac[i]
        @inbounds ctspo4[x[i], y[i], z[i]] += -plank.VPO4[i] * ΔT * ac[i]
    end
    return nothing
end

##### deal with grazed or dead individuals
function gpu_calc_loss_kernel!(ctsdoc, ctspoc, ctsdon, ctspon, ctsdop, ctspop, plank, ac, x, y, z,
                                   loss, lossFracC, lossFracN, lossFracP, R_NC, R_PC)
    index = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    stride = blockDim().x * gridDim().x
    for i = index:stride:size(ac,1)
        @inbounds CUDA.@atomic ctsdoc[x[i], y[i], z[i]] += (plank.Bm[i] + plank.Cq[i]) * lossFracC * ac[i] * loss[i]
        @inbounds CUDA.@atomic ctsdon[x[i], y[i], z[i]] += (plank.Bm[i]*R_NC + plank.Nq[i]) * lossFracN * ac[i] * loss[i]
        @inbounds CUDA.@atomic ctsdop[x[i], y[i], z[i]] += (plank.Bm[i]*R_PC + plank.Pq[i]) * lossFracP * ac[i] * loss[i]
        @inbounds CUDA.@atomic ctspoc[x[i], y[i], z[i]] += (plank.Bm[i] + plank.Cq[i]) * (1.0-lossFracC) * ac[i] * loss[i]
        @inbounds CUDA.@atomic ctspon[x[i], y[i], z[i]] += (plank.Bm[i]*R_NC + plank.Nq[i]) * (1.0-lossFracN) * ac[i] * loss[i]
        @inbounds CUDA.@atomic ctspop[x[i], y[i], z[i]] += (plank.Bm[i]*R_PC + plank.Pq[i]) * (1.0-lossFracP) * ac[i] * loss[i]
    end
    return nothing
end
function calc_loss!(ctsdoc, ctspoc, ctsdon, ctspon, ctsdop, ctspop, plank, ac, x, y, z,
                    loss, lossFracC, lossFracN, lossFracP, R_NC, R_PC, ::GPU)
    @cuda threads=256 blocks=ceil(Int, size(ac,1)/256) gpu_calc_loss_kernel!(ctsdoc, ctspoc, ctsdon, ctspon,
                                                        ctsdop, ctspop, plank, ac, x, y, z, loss,
                                                        lossFracC, lossFracN, lossFracP, R_NC, R_PC)
    return nothing
end
function calc_loss!(ctsdoc, ctspoc, ctsdon, ctspon, ctsdop, ctspop, plank, ac, x, y, z,
                           loss, lossFracC, lossFracN, lossFracP, R_NC, R_PC, ::CPU)
    for i in 1:size(ac,1)
        @inbounds ctsdoc[x[i], y[i], z[i]] += (plank.Bm[i] + plank.Cq[i]) * lossFracC * ac[i] * loss[i]
        @inbounds ctsdon[x[i], y[i], z[i]] += (plank.Bm[i]*R_NC + plank.Nq[i]) * lossFracN * ac[i] * loss[i]
        @inbounds ctsdop[x[i], y[i], z[i]] += (plank.Bm[i]*R_PC + plank.Pq[i]) * lossFracP * ac[i] * loss[i]
        @inbounds ctspoc[x[i], y[i], z[i]] += (plank.Bm[i] + plank.Cq[i]) * (1.0-lossFracC) * ac[i] * loss[i]
        @inbounds ctspon[x[i], y[i], z[i]] += (plank.Bm[i]*R_NC + plank.Nq[i]) * (1.0-lossFracN) * ac[i] * loss[i]
        @inbounds ctspop[x[i], y[i], z[i]] += (plank.Bm[i]*R_PC + plank.Pq[i]) * (1.0-lossFracP) * ac[i] * loss[i]
    end
    return nothing
end


using DrWatson
@quickactivate
using ArgParse
using GenerativeAD
import StatsBase: fit!, predict
using StatsBase
using BSON, FileIO
using Flux
using GenerativeModels
using DistributionsAD
using ValueHistories

s = ArgParseSettings()
@add_arg_table! s begin
	"modelname"
		default = "vae"
		arg_type = String
		help = "model name"
	"dataset"
		default = "iris"
		arg_type = String
		help = "dataset"
	"--seed"
		default = nothing
		help = "if specified, only results for a given seed will be recomputed"
	"--anomaly_class"
		default = nothing
		help = "if specified, only results for a given anomaly class will be recomputed"
end
parsed_args = parse_args(ARGS, s)
@unpack dataset, modelname, seed, anomaly_class = parsed_args

masterpath = datadir("experiments/images_leave-one-in/$(modelname)/$(dataset)")
files = GenerativeAD.Evaluation.collect_files(masterpath)
mfiles = filter(f->occursin("model", f), files)
if seed != nothing
	filter!(x->occursin("/seed=$seed/", x), mfiles)
end
if anomaly_class != nothing
	filter!(x->occursin("/ac=$(anomaly_class)/", x), mfiles)
end

sample_score_batched(m,x,L,batchsize) =
	vcat(map(y-> Base.invokelatest(GenerativeAD.Models.reconstruction_score, m, y, L),
		Flux.Data.DataLoader(x, batchsize=batchsize))...)
sample_score_batched_gpu(m,x,L,batchsize) =
	vcat(map(y-> cpu(Base.invokelatest(GenerativeAD.Models.reconstruction_score, m,
		gpu(Array(y)), L)), Flux.Data.DataLoader(x, batchsize=batchsize))...)

function save_sample_score(f::String, data, seed::Int, ac=nothing)
	# get model
	savepath = dirname(f)
	mdata = load(f)
	model = mdata["model"]
	L = 100

	# setup entries to be saved
	save_entries = (
		modelname = modelname,
		fit_t = mdata["fit_t"],
		history = mdata["history"],
		dataset = dataset,
		npars = sum(map(p->length(p), Flux.params(model))),
		model = nothing,
		seed = seed
		)
	save_entries = (ac == nothing) ? save_entries : merge(save_entries, (anomaly_class=ac,))
	result = (x -> sample_score_batched_gpu(gpu(model), x, L, 256),
		merge(mdata["parameters"], (L = L, score = "reconstruction-sampled")))

	# if the file does not exist already, compute the scores
	savef = joinpath(savepath, savename(result[2], "bson", digits=5))
	if !isfile(savef)
		@info "computing sample score for $f"
		GenerativeAD.experiment(result..., data, savepath; save_entries...)
	end
end

for f in mfiles
	# get data
	savepath = dirname(f)
	local seed = parse(Int, replace(basename(savepath), "seed=" => ""))
	ac = occursin("ac=", savepath) ? parse(Int, replace(basename(dirname(savepath)), "ac=" => "")) : nothing
	data = GenerativeAD.load_data(dataset, seed=seed, anomaly_class_ind=ac,
		method="leave-one-in")

	# compute and save the score
	save_sample_score(f, data, seed, ac)
end
@info "DONE"


function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:A2CGAE},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
    seed = 123,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_A2CGAE_CartPole_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)
    rng = StableRNG(seed)
    N_ENV = 16
    UPDATE_FREQ = 10
    env = MultiThreadEnv([
        CartPoleEnv(; T = Float32, rng = StableRNG(hash(seed + i))) for i in 1:N_ENV
    ])
    ns, na = length(get_state(env[1])), length(get_actions(env[1]))
    RLBase.reset!(env, is_force = true)
    agent = Agent(
        policy = QBasedPolicy(
            learner = A2CGAELearner(
                approximator = ActorCritic(
                    actor = NeuralNetworkApproximator(
                        model = Chain(
                            Dense(ns, 256, relu; initW = glorot_uniform(rng)),
                            Dense(256, na; initW = glorot_uniform(rng)),
                        ),
                        optimizer = ADAM(1e-3),
                    ),
                    critic = NeuralNetworkApproximator(
                        model = Chain(
                            Dense(ns, 256, relu; initW = glorot_uniform(rng)),
                            Dense(256, 1; initW = glorot_uniform(rng)),
                        ),
                        optimizer = ADAM(1e-3),
                    ),
                ) |> cpu,
                γ = 0.99f0,
                λ = 0.97f0,
                actor_loss_weight = 1.0f0,
                critic_loss_weight = 0.5f0,
                entropy_loss_weight = 0.001f0,
                update_freq = UPDATE_FREQ
            ),
            explorer = BatchExplorer(GumbelSoftmaxExplorer(;)),
        ),
        trajectory = CircularArraySARTTrajectory(;
            capacity = UPDATE_FREQ,
            state = Matrix{Float32} => (ns, N_ENV),
            action = Vector{Int} => (N_ENV,),
            reward = Vector{Float32} => (N_ENV,),
            terminal = Vector{Bool} => (N_ENV,),
        ),
    )
    stop_condition = StopAfterStep(50_000)
    total_reward_per_episode = TotalBatchRewardPerEpisode(N_ENV)
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env
            with_logger(lg) do
                @info(
                    "training",
                    actor_loss = agent.policy.learner.actor_loss,
                    critic_loss = agent.policy.learner.critic_loss,
                    entropy_loss = agent.policy.learner.entropy_loss,
                    loss = agent.policy.learner.loss,
                )
                for i in 1:length(env)
                    if get_terminal(env[i])
                        @info "training" reward = total_reward_per_episode.rewards[i][end] log_step_increment =
                            0
                        break
                    end
                end
            end
        end,
    )
    Experiment(
        agent,
        env,
        stop_condition,
        hook,
        "# A2CGAE with CartPole",
    )
end


##########################################################################
# The IndoorsNetwork is an indoors microcell network inspired by the ITU-R
# InH indoors network. The channel model is simplified; it is just using
# uncorrelated Rayleigh fading. The large-scale fading parameters should
# be close to the InH model though. The number of BSs has also been
# generalized compared to InH. The LoS is independent between BSs, for
# each MS. The same holds for the shadow fading. This is motivated by
# Sec. 3.3.1 in Deliverable 1.1.2 of IST-WINNER II.
#
# The typical transmit power is 21 dBm over 20 MHz (1200 subcarriers).
# This leads to 10*log10(10^(21/10)/1200) = -9.8 dBm for one 15 kHz subcarrier.

##########################################################################
# Network definition
type IndoorsNetwork{MS_t <: PhysicalMS, BS_t <: PhysicalBS, System_t <: System, PropagationEnvironment_t <: PropagationEnvironment} <: PhysicalNetwork
    MSs::Vector{MS_t}
    BSs::Vector{BS_t}

    system::System_t
    num_MSs_per_cell::Int
    propagation_environments::Dict{Symbol,PropagationEnvironment_t} # two keys only, :LoS and :NLoS
    corridor_length::Float64
    corridor_width::Float64
    guard_distance::Float64
    aux_network_params::AuxNetworkParams

    assignment::Assignment
end

# Convenience constructor without network params and assignments
IndoorsNetwork(MSs, BSs, system, num_MSs_per_cell, propagation_environments, corridor_length, corridor_width, guard_distance) =
    IndoorsNetwork(MSs, BSs, system, num_MSs_per_cell, propagation_environments, corridor_length, corridor_width, guard_distance, AuxNetworkParams(), Assignment())

Base.show(io::IO, x::IndoorsNetwork) =
    print(io, "Indoors(I = $(length(x.BSs)), Kc = $(x.num_MSs_per_cell), corridor_length = $(x.corridor_length), corridor_width = $(x.corridor_width), GD = $(x.guard_distance))")
Base.showcompact(io::IO, x::IndoorsNetwork) =
    print(io, "Indoors($(length(x.BSs)), $(x.num_MSs_per_cell), $(x.corridor_length), $(x.corridor_width), $(x.guard_distance))")

# The default parameter values are taken from ITU-R M.2135-1.
function setup_indoors_network(
    num_BSs, num_MSs_per_cell, num_MS_antennas, num_BS_antennas;
    system = SinglecarrierSystem(3.4e9, 15e3),
    propagation_environments = (@compat Dict(:LoS => SimpleLargescaleFadingEnvironment(16.9, 32.8 + 20*log10(3.4), 0, 3),
                                            :NLoS => SimpleLargescaleFadingEnvironment(43.3, 11.5 + 20*log10(3.4), 0, 4))),
    corridor_length = 120.,
    corridor_width = 50.,
    guard_distance = 3.,
    transmit_power = 10^(-9.8/10), transmit_powers = transmit_power*ones(Float64, num_BSs),
    BS_antenna_gain_params = [ OmnidirectionalAntennaParams(0) for i = 1:num_BSs ],
    user_priority = 1., user_priorities = ones(Float64, num_BSs*num_MSs_per_cell),
    num_streams = 1, num_streamss = num_streams*ones(Int, num_BSs*num_MSs_per_cell),
    MS_antenna_gain_dB = 0., MS_antenna_gains_dB = MS_antenna_gain_dB*ones(Float64, num_BSs*num_MSs_per_cell),
    receiver_noise_figure = 7., receiver_noise_figures = receiver_noise_figure*ones(Float64, num_BSs*num_MSs_per_cell))

    isa(num_MS_antennas, Vector) || (num_MS_antennas = num_MS_antennas*ones(Int, num_BSs*num_MSs_per_cell))
    isa(num_BS_antennas, Vector) || (num_BS_antennas = num_BS_antennas*ones(Int, num_BSs))

    # Place BSs
    y = corridor_width/2
    Δ = corridor_length/num_BSs
    BSs = Array(PhysicalBS, 0)
    for i = 1:num_BSs
        # BSs uniformly placed along corridor
        push!(BSs, PhysicalBS(num_BS_antennas[i], Position((i-1)*Δ + Δ/2, y), transmit_powers[i], BS_antenna_gain_params[i]))
    end

    MSs = [ PhysicalMS(num_MS_antennas[k], Position(0, 0), Velocity(0, 0), user_priorities[k], num_streamss[k], MS_antenna_gains_dB[k], receiver_noise_figures[k], SimpleLargescaleFadingEnvironmentState(zeros(Float64, num_BSs), falses(num_BSs))) for k = 1:num_BSs*num_MSs_per_cell ]

    IndoorsNetwork(MSs, BSs, system, num_MSs_per_cell,
        propagation_environments, corridor_length, corridor_width, guard_distance)
end

##########################################################################
# Standard cell assignment functions
function IDCellAssignment!(channel, network::IndoorsNetwork)
    Kc = network.num_MSs_per_cell; I = get_num_BSs(network)
    cell_assignment = Array(Int, I*Kc)

    for i = 1:I
        cell_assignment[(i-1)*Kc+1:i*Kc] = i
    end

    network.assignment = Assignment(cell_assignment, I)

    return AssignmentResults()
end

# Do this for now, even though the shadowing might actually give a certain MS
# better connection to a BS in another cell.
LargeScaleFadingCellAssignment!(channel, network::IndoorsNetwork) =
    IDCellAssignment!(channel, network)

##########################################################################
# Simulation functions
function draw_user_drop!{MS_t <: PhysicalMS, BS_t <: PhysicalBS, System_t <: System}(network::IndoorsNetwork{MS_t, BS_t, System_t, SimpleLargescaleFadingEnvironment})
    I = get_num_BSs(network); K = get_num_MSs(network)
    Δ = network.corridor_length/I

    # ITU-R M.2135-1, p. 33
    P_LoS = d -> begin
        if d <= 18
            return 1.
        elseif d < 37
            return exp(-(d - 18)/27)
        else
            return 0.5
        end
    end

    for k = 1:K
        i = div(k - 1, network.num_MSs_per_cell) + 1 # serving BS id

        # MS position
        while true
            x = (i-1)*Δ + Δ*rand()
            y = network.corridor_width*rand()
            candidate_position = Position(x, y)

            if get_distance(network.BSs[i].position, candidate_position) > network.guard_distance
                network.MSs[k].position = candidate_position
                break
            end
        end

        # Shadow fading and LoS/NLoS. Uncorrelated between BS links.
        shadow_realizations_dB = zeros(Float64, I)
        LoSs = BitArray(I)
        for j = 1:I
            if rand() <= P_LoS(get_distance(network.BSs[j].position, network.MSs[k].position))
                LoSs[j] = true
                shadow_realizations_dB[j] = network.propagation_environments[:LoS].shadow_sigma_dB*randn()
            else
                LoSs[j] = false
                shadow_realizations_dB[j] = network.propagation_environments[:NLoS].shadow_sigma_dB*randn()
            end
        end
        network.MSs[k].propagation_environment_state =
            SimpleLargescaleFadingEnvironmentState(shadow_realizations_dB, LoSs)
    end
end

function draw_channel{MS_t <: PhysicalMS, BS_t <: PhysicalBS}(network::IndoorsNetwork{MS_t, BS_t, SinglecarrierSystem, SimpleLargescaleFadingEnvironment})
    K = get_num_MSs(network); I = get_num_BSs(network)
    Ns = get_num_MS_antennas(network); Ms = get_num_BS_antennas(network)

    coefs = Array(Matrix{Complex128}, K, I)
    large_scale_fading_factor = Array(Float64, K, I)

    distances = get_distances(network)

    for k = 1:K
        for i = 1:I
            # Small scale fading
            coefs[k,i] = (1/sqrt(2))*(randn(Ns[k], Ms[i])
                                 + im*randn(Ns[k], Ms[i]))

            # Pathloss
            if network.MSs[k].propagation_environment_state.LoSs[i]
                pathloss_factor =
                    sqrt(10^(-(network.propagation_environments[:LoS].pathloss_beta + network.propagation_environments[:LoS].pathloss_alpha*log10(distances[k,i]))/10))
            else
                pathloss_factor =
                    sqrt(10^(-(network.propagation_environments[:NLoS].pathloss_beta + network.propagation_environments[:NLoS].pathloss_alpha*log10(distances[k,i]))/10))
            end

            # Shadow fading
            shadow_factor = sqrt(10^(network.MSs[k].propagation_environment_state.shadow_realizations_dB[i]/10))

            # BS antenna gain
            bs_antenna_gain = sqrt(get_antenna_gain(network.BSs[i].antenna_params))

            # MS antenna gain
            ms_antenna_gain = sqrt(10^(network.MSs[k].antenna_gain_dB/10))

            # Apply scale factors
            large_scale_fading_factor[k,i] = pathloss_factor*shadow_factor*bs_antenna_gain*ms_antenna_gain
            coefs[k,i] *= large_scale_fading_factor[k,i]
        end
    end

    return SinglecarrierChannel(coefs, Ns, Ms, K, I, large_scale_fading_factor)
end

##########################################################################
# Visualization functions
function plot_network_layout(network::IndoorsNetwork)
    I = get_num_BSs(network); K = get_num_MSs(network)
    Δ = network.corridor_length/I
    δ = network.corridor_width

    fig = PyPlot.figure()
    ax = fig[:add_subplot](1, 1, 1)

    # Corridor outline
    pos1 = network.BSs[1].position; pos2 = network.BSs[end].position
    sw = (pos1.x - Δ/2, pos1.y - δ/2); se = (pos2.x + Δ/2, pos2.y - δ/2)
    nw = (pos1.x - Δ/2, pos1.y + δ/2); ne = (pos2.x + Δ/2, pos2.y + δ/2)
    ax[:plot]([sw[1], se[1]], [sw[2], se[2]]; color="b", linestyle="-", linewidth=4)
    ax[:plot]([se[1], ne[1]], [se[2], ne[2]]; color="b", linestyle="-", linewidth=4)
    ax[:plot]([ne[1], nw[1]], [ne[2], nw[2]]; color="b", linestyle="-", linewidth=4)
    ax[:plot]([nw[1], sw[1]], [nw[2], sw[2]]; color="b", linestyle="-", linewidth=4)

    # BSs
    for i = 1:I
        pos = network.BSs[i].position
        ax[:plot](pos.x, pos.y; marker="x", color="b", markersize=10)

        # Cell outline
        if i < I
            ax[:plot]([pos.x + Δ/2, pos.x + Δ/2], [pos.y - δ/2, pos.y + δ/2]; color="b", linestyle="--", linewidth=2)
        end
    end

    # MSs
    for k = 1:K
        pos = network.MSs[k].position
        ax[:plot](pos.x, pos.y; marker="o", color="r", markersize=10)
    end

    display(fig)
end


export Geometry

function Geometry(con::H5SeisCon)
  if isnothing(con.seis)
    @error "Seis is Nothing"
    return
  end

  if isnothing(con.pkey) || isnothing(con.xkey) || isnothing(con.ykey) || isnothing(con.zkey)
    @error "pkey/xkey/ykey/zkey is Nothing"
    return
  end

  if isnothing(con.pkeyvals)
    @error "pkeyvals is Nothing"
    return
  end

  dt = Float32(abs(con.seis.getSampRate("ms")))
  nt = Integer(con.seis.getNSamp())
  t = Float32((nt-1)*dt)

  nsrc = length(con.pkeyvals)
  xCell = Vector{Vector{Float32}}(undef, nsrc)
  yCell = Vector{Vector{Float32}}(undef, nsrc)
  zCell = Vector{Vector{Float32}}(undef, nsrc)
  dtCell = Vector{Float32}(undef, nsrc)
  ntCell = Vector{Integer}(undef, nsrc)
  tCell = Vector{Float32}(undef, nsrc)
  keylist = [con.pkey]
  for block in 1:nsrc
    minlist = [con.pkeyvals[block]]
    maxlist = [con.pkeyvals[block]]
    _, _, ind = con.seis.getSortedData(keylist, minlist, maxlist, 0, 0)
    xy = Float32.(con.seis.getXYTraceHeaders([con.xkey, con.ykey], ind, "m", false))  # MUST BE TRUE

    xCell[block] = xy[:,1]
    yCell[block] = xy[:,2]
    zCell[block] = Float32.(con.seis.getTraceHeader([con.zkey], ind, [con.seis.getLengthUnits()], ["m"])[:])
    dtCell[block] = dt
    ntCell[block] = nt
    tCell[block] = t
  end

  return GeometryIC{Float32}(xCell, yCell, zCell, dtCell, ntCell, tCell)
end


module Remove_response_obspy
export remove_response_obspy, remove_response_obspy!
using PyCall, SeisIO

# Suppressing warning associated with PyCall
using Suppressor

"""
remove_response_obspy(S::SeisDatat)
remove Instrumental response with obspy.
Currently available only one channel.
This function will be replaced by julia original function
"""
function remove_response_obspy!(S::SeisData, xf::String; pre_filt::NTuple{4,Float64}=(0.001, 0.002, 10.0, 20.0), output::String="VEL")

    @suppress begin
    Stream          = pyimport("obspy.core.stream")["Stream"]
    Trace           = pyimport("obspy.core.trace")["Trace"]
    read_inventory  = pyimport("obspy")["read_inventory"]
    UTCDateTime     = pyimport("obspy.core.utcdatetime")["UTCDateTime"]

    trace=Trace()
    trace.data = S.x[1]
    trace.stats.sampling_rate = S.fs
    trace.stats.delta = 1.0./ S.fs
    trace.stats.starttime = UTCDateTime(string(u2d( S.t[1][1,2]*1e-6)))
    str1 = split(S.id[1], ".")
    trace.stats.network = str1[1]
    trace.stats.station=  str1[2]
    trace.stats.location= str1[3]
    trace.stats.channel=  str1[4]

    stream = Stream(traces=trace)
    #read stationXML
    inv = read_inventory(xf)
    #stresm.remove_sensitivity(inventory=inv)
    stream.remove_response(inventory=inv, output=output, pre_filt=pre_filt, plot=false)
    S.x[1] = stream.traces[1].data
    #add note
    SeisIO.note!(S, "remove_response_obspy!, pre_filt=$pre_filt")
    end
    return nothing
end

function remove_response_obspy(S::SeisData; pre_filt::NTuple{4,Float64}=(0.001, 0.002, 10.0, 20.0), output::String="VEL")
  U = deepcopy(S)
  remove_response_obspy!(U, pre_filt=pre_filt, output=output)
  return U
end

end


module LinearAlgebra

using Reexport

include("GLA.jl")
include("StructuredGLA.jl")

@reexport using .GraphicalLinearAlgebra
@reexport using .StructuredGraphicalLinearAlgebra

end


"""
    append_hs!(model::JuMP.Model, x, nh::Int)

Append equalty constraints (auto-generated)

# Arguments
    `model::JuMP.Model`: optimization model
    `x::JuMP.variables`: optimization model variables
    `nh::Int`: number of equality constraints
"""
function append_hs!(model::JuMP.Model, x, nh::Int)
    for ih = 1:nh
        if ih == 1
            @NLconstraint(model, h1(x...) == 0)
        elseif ih == 2
            @NLconstraint(model, h2(x...) == 0)
        elseif ih == 3
            @NLconstraint(model, h3(x...) == 0)
        elseif ih == 4
            @NLconstraint(model, h4(x...) == 0)
        elseif ih == 5
            @NLconstraint(model, h5(x...) == 0)
        elseif ih == 6
            @NLconstraint(model, h6(x...) == 0)
        elseif ih == 7
            @NLconstraint(model, h7(x...) == 0)
        elseif ih == 8
            @NLconstraint(model, h8(x...) == 0)
        elseif ih == 9
            @NLconstraint(model, h9(x...) == 0)
        elseif ih == 10
            @NLconstraint(model, h10(x...) == 0)
        elseif ih == 11
            @NLconstraint(model, h11(x...) == 0)
        elseif ih == 12
            @NLconstraint(model, h12(x...) == 0)
        elseif ih == 13
            @NLconstraint(model, h13(x...) == 0)
        elseif ih == 14
            @NLconstraint(model, h14(x...) == 0)
        elseif ih == 15
            @NLconstraint(model, h15(x...) == 0)
        elseif ih == 16
            @NLconstraint(model, h16(x...) == 0)
        elseif ih == 17
            @NLconstraint(model, h17(x...) == 0)
        elseif ih == 18
            @NLconstraint(model, h18(x...) == 0)
        elseif ih == 19
            @NLconstraint(model, h19(x...) == 0)
        elseif ih == 20
            @NLconstraint(model, h20(x...) == 0)
        end
    end
end


# %% imports

include("../../../Modules/CMAES/cmaes.jl")
include("EI_sim.jl")
include("../../../Modules/save_param.jl")

# %%

function optim()
    cfg = YAML.load(open("Julia/CartPole/LIF_base/cfg_cmaes.yml"))
    arch = cfg["arch"]

    N_alpha = arch[1]*4
    N_w = sum([arch[i]*(arch[i-1]-(div(arch[i-1],4))) for i in 2:length(arch)]) +1
    N = N_alpha + N_w

    best = Array{Float64}[]
    for i in 1:20
        push!(best,[])
    end
    best_fit = [-1.0 for i in 1:20]
    error = false

    # Create initial point (optional)
    results = load_param("Julia/CartPole/LIF_base/results_EI888.jld")
    genes_init = results["weights"][19]

    lambda = 4+floor(Int64,3*log(N))
    mu = floor(Int64,lambda/2)
    c = CMAES(genes_init, N=N, μ=mu, λ=lambda , τ=sqrt(N), τ_c=N^2, τ_σ=sqrt(N))
    for i in 1:500
        @time step!(c)
        bestind = argmin(c.F_λ)
        maxfit = -c.F_λ[bestind]
        println(i," ", maxfit, " (current best = $(maximum(best_fit)))")
        if maxfit > best_fit[1]
            best_fit[1] = maxfit
            best[1] = copy(c.offspring[bestind])
            idx = sortperm(best_fit)
            best_fit = best_fit[idx]
            best = best[idx]
        end
        if sum(best_fit) >= 4000
            break
        end
        if i > 20 && best_fit[end] < 190
            error = true
            break
        end
        if length(best[1]) > 0 && i%20 == 0
            save_param(best_fit,best)
        end
    end
    if best_fit[end] < 195
        error = true
    end
    return best_fit, best, error
end

# %% run sim

for opt in 1:20
    best_fit, best, error = optim()
    if error==false
        break
    end
end


include("dependencies_for_runtests.jl")

using Oceananigans.Grids: topology, XRegLatLonGrid, YRegLatLonGrid, ZRegLatLonGrid

function get_fields_from_checkpoint(filename)
    file = jldopen(filename)

    tracers = keys(file["tracers"])
    tracers = Tuple(Symbol(c) for c in tracers)

    velocity_fields = (u = file["velocities/u/data"],
                       v = file["velocities/v/data"],
                       w = file["velocities/w/data"])

    tracer_fields =
        NamedTuple{tracers}(Tuple(file["tracers/$c/data"] for c in tracers))

    current_tendency_velocity_fields = (u = file["timestepper/Gⁿ/u/data"],
                                        v = file["timestepper/Gⁿ/v/data"],
                                        w = file["timestepper/Gⁿ/w/data"])

    current_tendency_tracer_fields =
        NamedTuple{tracers}(Tuple(file["timestepper/Gⁿ/$c/data"] for c in tracers))

    previous_tendency_velocity_fields = (u = file["timestepper/G⁻/u/data"],
                                         v = file["timestepper/G⁻/v/data"],
                                         w = file["timestepper/G⁻/w/data"])

    previous_tendency_tracer_fields =
        NamedTuple{tracers}(Tuple(file["timestepper/G⁻/$c/data"] for c in tracers))

    close(file)

    solution = merge(velocity_fields, tracer_fields)
    Gⁿ = merge(current_tendency_velocity_fields, current_tendency_tracer_fields)
    G⁻ = merge(previous_tendency_velocity_fields, previous_tendency_tracer_fields)

    return solution, Gⁿ, G⁻
end

include("regression_tests/thermal_bubble_regression_test.jl")
include("regression_tests/rayleigh_benard_regression_test.jl")
include("regression_tests/ocean_large_eddy_simulation_regression_test.jl")

@testset "Nonhydrostatic Regression" begin
    @info "Running nonhydrostatic regression tests..."

    for arch in archs
        for grid_type in [:regular, :vertically_unstretched]
            @testset "Thermal bubble [$(typeof(arch)), $grid_type grid]" begin
                @info "  Testing thermal bubble regression [$(typeof(arch)), $grid_type grid]"
                run_thermal_bubble_regression_test(arch, grid_type)
            end

            @testset "Rayleigh–Bénard tracer [$(typeof(arch)), $grid_type grid]]" begin
                @info "  Testing Rayleigh–Bénard tracer regression [$(typeof(arch)), $grid_type grid]"
                run_rayleigh_benard_regression_test(arch, grid_type)
            end

            for closure in (AnisotropicMinimumDissipation(ν=1.05e-6, κ=1.46e-7), SmagorinskyLilly(C=0.23, Cb=1, Pr=1, ν=1.05e-6, κ=1.46e-7))
                closurename = string(typeof(closure).name.wrapper)
                @testset "Ocean large eddy simulation [$(typeof(arch)), $closurename, $grid_type grid]" begin
                    @info "  Testing oceanic large eddy simulation regression [$(typeof(arch)), $closurename, $grid_type grid]"
                    run_ocean_large_eddy_simulation_regression_test(arch, grid_type, closure)
                end
            end
        end
    end
end


using GLVisualize, Colors, GeometryTypes, Reactive, GLAbstraction, GLWindow

if !isdefined(:runtests)
    window = glscreen()
end

description = """
Simple example showing how to create color choosers.
You can drag with right and left mouse click on the color to change it:
↔ right → red   | ↔ left → blue
↕ right → green | ↕ left → alpha
"""

color_a = Signal(RGBA{Float32}(1,0,0,0.7))
a_v, color_s = widget(color_a, window)

color_b = Signal(RGBA{Float32}(0,1,0,0.7))
b_v, color_s = widget(color_b, window)

color_c = Signal(RGBA{Float32}(0,0,1,0.7))
c_v, color_s = widget(color_c, window)
edit_screen = Screen(window, area=map(window.area) do a
    SimpleRectangle(0,0,60,a.h)
end)
_view(layout!(SimpleRectangle{Float32}(0,0,60,60), a_v), edit_screen, camera=:fixed_pixel)
_view(layout!(SimpleRectangle{Float32}(0,60,60,60), b_v), edit_screen, camera=:fixed_pixel)
_view(layout!(SimpleRectangle{Float32}(0,120,60,60), c_v), edit_screen, camera=:fixed_pixel)

function multirandomwalk(n1, n2)
    a = rand(Point3f0, n1)*2f0
    r = Array{Point3f0}(n1*n2)
    i = 1
    for j=1:n1
        pstart = a[j]
        for k=1:n2
            pstart += (rand(Point3f0)*2f0 - 1f0)*0.07f0
            r[i] = pstart
            i+=1
        end
    end
    r
end

a = multirandomwalk(15, 20)
b = multirandomwalk(15, 20)
c = multirandomwalk(15, 20)

_view(visualize(
    (Circle(Point2f0(0), 0.02f0), a),
    color=color_a, billboard=true
), window, camera=:perspective)

_view(visualize(
    (Circle(Point2f0(0), 0.02f0), b),
    color=color_b, billboard=true
), window, camera=:perspective)

_view(visualize(
    (Circle(Point2f0(0), 0.02f0), c),
    color=color_c, billboard=true
), window, camera=:perspective)

if !isdefined(:runtests)
    renderloop(window)
end