import openmdao.api as om
import aviary.api as av


subsystem_options = {'core_aerodynamics':
                     {'method': 'low_speed',
                      'ground_altitude': 0.,  # units='ft'
                      'angles_of_attack': [
                          -5.0, -4.0, -3.0, -2.0, -1.0,
                          0.0, 1.0, 2.0, 3.0, 4.0, 5.0,
                          6.0, 7.0, 8.0, 9.0, 10.0, 11.0,
                          12.0, 13.0, 14.0, 15.0],  # units='deg'
                      'lift_coefficients': [
                          0.01, 0.1, 0.2, 0.3, 0.4,
                          0.5178, 0.6, 0.75, 0.85, 0.95, 1.05,
                          1.15, 1.25, 1.35, 1.5, 1.6, 1.7,
                          1.8, 1.85, 1.9, 1.95],
                      'drag_coefficients': [
                          0.04, 0.02, 0.01, 0.02, 0.04,
                          0.0674, 0.065, 0.065, 0.07, 0.072, 0.076,
                          0.084, 0.09, 0.10, 0.11, 0.12, 0.13,
                          0.15, 0.16, 0.18, 0.20],
                      'lift_coefficient_factor': 1.,
                      'drag_coefficient_factor': 1.}}

optimize_mach = True
optimize_altitude = True
optimizer = "SLSQP"

phase_info = {
    "pre_mission": {"include_takeoff": False, "optimize_mass": False},
    'AB': {
        'user_options': {
            'num_segments': 5,
            'order': 3,
            'fix_initial': True,
            'ground_roll': True,
            'duration_ref': (100., 'kn'),
            'duration_bounds': ((100., 500.), 'kn'),
        },
        'subsystem_options': subsystem_options,
        'initial_guesses': {
            'distance': [(0., 2.e3), 'ft'],
            'time': [(0., 20.), 's'],
            'velocity': [(1., 120.), 'kn'],
            'mass': [(175.e3, 174.85e3), 'lbm'],
        },
    },
    'rotate': {
        'user_options': {
            'num_segments': 5,
            'order': 3,
            'fix_initial': False,
            'ground_roll': True,
            'clean': False,
            'initial_ref': (1.e3, 'ft'),
            'initial_bounds': ((1.e3, 3.e3), 'ft'),
            'duration_ref': (1.e3, 'ft'),
            'duration_bounds': ((200., 2.e3), 'ft'),
            'mach_bounds': ((0.18, 0.2), 'unitless'),
            'polynomial_control_order': 1,
            'throttle_enforcement': 'boundary_constraint',
            'optimize_mach': optimize_mach,
            'optimize_altitude': False,
            'rotation': True,
            'initial_mach': (0.18, 'unitless'),
            'final_mach': (0.2, 'unitless'),
            'initial_altitude': (0., 'ft'),
            'final_altitude': (0., 'ft'),
            'constraints': {
                'normal_force': {
                    'equals': 0.,
                    'loc': 'final',
                    'units': 'lbf',
                    'type': 'boundary',
                    'ref': 10.e5,
                },
            },
        },
        'subsystem_options': subsystem_options,
        'initial_guesses': {
            'distance': [(2.e3, 1.e3), 'ft'],
            'time': [(20., 25.), 's'],
            'mass': [(174.85e3, 174.84e3), 'lbm'],
            'alpha': [(0., 12.), 'deg'],
        },
    },
    'BC': {
        'user_options': {
            'num_segments': 5,
            'order': 3,
            'fix_initial': False,
            'ground_roll': False,
            'clean': False,
            'initial_ref': (1.e3, 'ft'),
            'initial_bounds': ((1., 16.e3), 'ft'),
            'duration_ref': (1.e3, 'ft'),
            'duration_bounds': ((500., 1500.), 'ft'),
            'mach_bounds': ((0.2, 0.22), 'unitless'),
            'altitude_bounds': ((0., 250.), 'ft'),
            'initial_mach': (0.2, 'unitless'),
            'final_mach': (0.22, 'unitless'),
            'initial_altitude': (0., 'ft'),
            'final_altitude': (50., 'ft'),
            'polynomial_control_order': 1,
            'throttle_enforcement': 'boundary_constraint',
            'optimize_mach': optimize_mach,
            'optimize_altitude': optimize_altitude,
            'rotation': False,
        },
        'subsystem_options': subsystem_options,
        'initial_guesses': {
            'distance': [(3.e3, 1.e3), 'ft'],
            'time': [(25., 35.), 's'],
            'mass': [(174.84e3, 174.82e3), 'lbm'],
        },
    },
    'CD_to_P2': {
        'user_options': {
            'num_segments': 4,
            'order': 3,
            'fix_initial': False,
            'ground_roll': False,
            'clean': False,
            'initial_ref': (1.e3, 'ft'),
            'initial_bounds': ((1.e3, 20.e3), 'ft'),
            'duration_ref': (1.e3, 'ft'),
            'duration_bounds': ((3.e3, 20.e3), 'ft'),
            'mach_bounds': ((0.22, 0.3), 'unitless'),
            'altitude_bounds': ((0., 985.), 'ft'),
            'initial_mach': (0.22, 'unitless'),
            'final_mach': (0.3, 'unitless'),
            'initial_altitude': (50., 'ft'),
            'final_altitude': (985., 'ft'),
            'polynomial_control_order': 1,
            'throttle_enforcement': 'boundary_constraint',
            'optimize_mach': optimize_mach,
            'optimize_altitude': optimize_altitude,
            'constraints': {
                'altitude': {
                    'equals': 985.,
                    'loc': 'final',
                    'units': 'ft',
                    'type': 'boundary',
                },
            },
        },
        'subsystem_options': subsystem_options,
        'initial_guesses': {
            'distance': [(4.e3, 10.e3), 'ft'],
            'time': [(35., 60.), 's'],
            'mass': [(174.82e3, 174.8e3), 'lbm'],
        },
    },
    'P2_to_DE': {
        'user_options': {
            'num_segments': 4,
            'order': 3,
            'fix_initial': False,
            'ground_roll': False,
            'clean': False,
            'initial_ref': (1.e3, 'ft'),
            'initial_bounds': ((1.e3, 20.e3), 'ft'),
            'duration_ref': (1.e3, 'ft'),
            'duration_bounds': ((3.e3, 20.e3), 'ft'),
            'mach_bounds': ((0.22, 0.3), 'unitless'),
            'altitude_bounds': ((985., 1100.), 'ft'),
            'initial_mach': (0.22, 'unitless'),
            'final_mach': (0.3, 'unitless'),
            'initial_altitude': (985., 'ft'),
            'final_altitude': (1100., 'ft'),
            'polynomial_control_order': 1,
            'throttle_enforcement': 'path_constraint',
            'optimize_mach': optimize_mach,
            'optimize_altitude': optimize_altitude,
            'constraints': {
                'distance': {
                    'upper': 19.e3,
                    'ref': 20.e3,
                    'loc': 'final',
                    'units': 'ft',
                    'type': 'boundary',
                },
            },
        },
        'subsystem_options': subsystem_options,
        'initial_guesses': {
            'distance': [(10.e3, 14.e3), 'ft'],
            'time': [(60., 80.), 's'],
            'mass': [(174.8e3, 174.5e3), 'lbm'],
        },
    },
    'DE': {
        'user_options': {
            'num_segments': 3,
            'order': 3,
            'fix_initial': False,
            'ground_roll': False,
            'clean': False,
            'initial_ref': (1.e3, 'ft'),
            'initial_bounds': ((500., 30.e3), 'ft'),
            'duration_ref': (1.e3, 'ft'),
            'duration_bounds': ((50., 5000.), 'ft'),
            'mach_bounds': ((0.24, 0.32), 'unitless'),
            'altitude_bounds': ((985., 1.5e3), 'ft'),
            'initial_mach': (0.3, 'unitless'),
            'final_mach': (0.3, 'unitless'),
            'initial_altitude': (1100., 'ft'),
            'final_altitude': (1200., 'ft'),
            'polynomial_control_order': 2,
            'optimize_mach': optimize_mach,
            'optimize_altitude': optimize_altitude,
            'throttle_enforcement': 'path_constraint',
            'constraints': {
                'flight_path_angle': {
                    'equals': 4.,
                    'loc': 'final',
                    'units': 'deg',
                    'type': 'boundary',
                },
            },
        },
        'subsystem_options': subsystem_options,
        'initial_guesses': {
            'distance': [(18.e3, 2.e3), 'ft'],
            'mass': [(174.5e3, 174.4e3), 'lbm'],
            'time': [(80., 85.), 's'],
        },
    },
    'EF_to_P1': {
        'user_options': {
            'num_segments': 3,
            'order': 3,
            'fix_initial': False,
            'ground_roll': False,
            'clean': False,
            'initial_ref': (1.e3, 'ft'),
            'initial_bounds': ((500., 50.e3), 'ft'),
            'duration_ref': (1.e3, 'ft'),
            'duration_bounds': ((1.e3, 20.e3), 'ft'),
            'mach_bounds': ((0.24, 0.32), 'unitless'),
            'altitude_bounds': ((1.1e3, 1.2e3), 'ft'),
            'initial_mach': (0.3, 'unitless'),
            'final_mach': (0.3, 'unitless'),
            'initial_altitude': (1100., 'ft'),
            'final_altitude': (1200., 'ft'),
            'polynomial_control_order': 1,
            'optimize_mach': optimize_mach,
            'optimize_altitude': optimize_altitude,
            'throttle_enforcement': 'bounded',
            'constraints': {
                'distance': {
                    'equals': 21325.,
                    'units': 'ft',
                    'type': 'boundary',
                    'loc': 'final',
                    'ref': 30.e3,
                },
                'flight_path_angle': {
                    'equals': 4.,
                    'loc': 'final',
                    'units': 'deg',
                    'type': 'boundary',
                },
            },
        },
        'subsystem_options': subsystem_options,
        'initial_guesses': {
            'distance': [(20.e3, 1325.), 'ft'],
            'mass': [(174.4e3, 174.3e3), 'lbm'],
            'time': [(85., 90.), 's'],
        },
    },
    'EF_past_P1': {
        'user_options': {
            'num_segments': 5,
            'order': 3,
            'fix_initial': False,
            'ground_roll': False,
            'clean': False,
            'initial_ref': (1.e3, 'ft'),
            'initial_bounds': ((20.e3, 50.e3), 'ft'),
            'duration_ref': (1.e3, 'ft'),
            'duration_bounds': ((100., 50.e3), 'ft'),
            'mach_bounds': ((0.24, 0.32), 'unitless'),
            'altitude_bounds': ((1.e3, 3.e3), 'ft'),
            'initial_mach': (0.3, 'unitless'),
            'final_mach': (0.3, 'unitless'),
            'initial_altitude': (1200., 'ft'),
            'final_altitude': (2000., 'ft'),
            'polynomial_control_order': 1,
            'optimize_mach': optimize_mach,
            'optimize_altitude': optimize_altitude,
            'throttle_enforcement': 'boundary_constraint',
            'constraints': {
                'flight_path_angle': {
                    'equals': 4.,
                    'loc': 'final',
                    'units': 'deg',
                    'type': 'boundary',
                },
                'distance': {
                    'equals': 30.e3,
                    'units': 'ft',
                    'type': 'boundary',
                    'loc': 'final',
                    'ref': 30.e3,
                },
            },
        },
        'subsystem_options': subsystem_options,
        'initial_guesses': {
            'distance': [(21325., 50.e3), 'ft'],
            'mass': [(174.3e3, 174.2e3), 'lbm'],
            'time': [(90., 180.), 's'],
        },
    },
    "post_mission": {
        "include_landing": False,
        "constrain_range": False,
    },
}

if __name__ == '__main__':
    prob = av.AviaryProblem()

    # Load aircraft and options data from user
    # Allow for user overrides here
    prob.load_inputs(
        'models/test_aircraft/aircraft_for_bench_solved2dof.csv', phase_info)

    # Preprocess inputs
    prob.check_and_preprocess_inputs()

    prob.add_pre_mission_systems()

    prob.add_phases()

    prob.add_post_mission_systems()

    # Link phases and variables
    prob.link_phases()

    prob.add_driver(optimizer, max_iter=25)

    prob.add_design_variables()

    # Load optimization problem formulation
    # Detail which variables the optimizer can control
    prob.add_objective('mass')

    prob.setup()

    prob.set_initial_guesses()

    prob.run_aviary_problem(record_filename='detailed_takeoff.db')

    try:
        loc = prob.get_outputs_dir()
        cr = om.CaseReader(f'{loc}/detailed_takeoff.db')
    except:
        cr = om.CaseReader('detailed_takeoff.db')

    cases = cr.get_cases('problem')
    case = cases[0]

    output_data = {}

    for (point_name, phase_name) in [['P1', 'EF_to_P1'], ['P2', 'CD_to_P2']]:
        output_data[point_name] = {}
        output_data[point_name]['thrust_fraction'] = case.get_val(f'traj.{phase_name}.rhs_all.thrust_net', units='N')[
            -1][0] / case.get_val(f'traj.{phase_name}.rhs_all.thrust_net_max', units='N')[-1][0]
        output_data[point_name]['true_airspeed'] = case.get_val(
            f'traj.{phase_name}.timeseries.velocity', units='kn')[-1][0]
        output_data[point_name]['angle_of_attack'] = case.get_val(
            f'traj.{phase_name}.timeseries.alpha', units='deg')[-1][0]
        output_data[point_name]['flight_path_angle'] = case.get_val(
            f'traj.{phase_name}.timeseries.flight_path_angle', units='deg')[-1][0]
        output_data[point_name]['altitude'] = case.get_val(
            f'traj.{phase_name}.timeseries.altitude', units='ft')[-1][0]
        output_data[point_name]['distance'] = case.get_val(
            f'traj.{phase_name}.timeseries.distance', units='ft')[-1][0]

    print(output_data)
