trait EngineLosses:
    engine_loss_shaft is DriveTrain.Shaft
    tc_turbine_loss_shaft is DriveTrain.Shaft
    gear_output_loss_shaft is DriveTrain.Shaft

    engine_losses is DriveTrain.ViscousGear:
        ratio: 1.0
        connectors: [engine_shaft.output,engine_loss_shaft.input]

    # Viscous loss between torque converter and gear box
    tc_to_gear_viscous is DriveTrain.ViscousGear:
        ratio: 1.0
        connectors: [torque_converter_shaft.output, tc_turbine_loss_shaft.input]

    # Viscous loss between gearbox output and prop shaft / center split
    gear_to_drive_viscous is DriveTrain.ViscousGear:
        ratio: 1.0
        connectors: [gear_shaft.output, gear_output_loss_shaft.input]

    # Reroute drivetrain connectors so the new lossy components are included
    torque_converter:
        connectors: [engine_loss_shaft.output, torque_converter_shaft.input]
    gear_box:
        connectors: [tc_turbine_loss_shaft.output, gear_shaft.input]
    mid_center_axle:
        drive_shaft: gear_output_loss_shaft
        left_axle_shaft: front_drive_shaft
        right_axle_shaft: midrear_split_shaft

DriveTrainKomatsuHM400WithEngineLoss is DrivetrainKomatsuHM400:
    with EngineLosses

    engine_loss_shaft.inertia.inertia:      0.1
    tc_turbine_loss_shaft.inertia.inertia:  0.1
    gear_output_loss_shaft.inertia.inertia: 0.1

    engine_losses:
        dissipation.viscosity:  600
    tc_to_gear_viscous:
        dissipation.viscosity:  60000
    gear_to_drive_viscous:
        dissipation.viscosity:  80000

DrivetrainKomatsuHM400 is SixWheelTemplate:
    .doc: """
        Engine model: Komatsu SAA6D140E-7

        Based on data from: 
        [1]   Komatsu HM400-5 - articulated dump truck
              https://www.komatsu.com/content/dam/komatsu/websites/en-us/miscellaneous-web-site-assets/HM400_5_AESS878_0225_v2.pdf

        [2]   Komatsu HD405-6 - fixed frame truck that uses the same engine size but an older generation
              https://www.komatsu.eu/Assets/GetBrochureByProductName.aspx?id=HD325-6&langID=en

        [3]   Komatsu HM400-5 Shop Manual (SEN06519-06)
    """
    # Main drivetrain shafts
    engine_shaft.inertia.inertia:               4.000
    torque_converter_shaft.inertia.inertia:     0.300
    gear_shaft.inertia.inertia:                 0.250
    midrear_split_shaft.inertia.inertia:        0.080

    # Differential input shafts
    front_drive_shaft.inertia.inertia:          0.1
    mid_drive_shaft.inertia.inertia:            0.1
    rear_drive_shaft.inertia.inertia:           0.1

    # Differential output shafts
    front_left_half_shaft.inertia.inertia:      0.1
    front_right_half_shaft.inertia.inertia:     0.1
    mid_left_half_shaft.inertia.inertia:        0.1
    mid_right_half_shaft.inertia.inertia:       0.1
    rear_left_half_shaft.inertia.inertia:       0.1
    rear_right_half_shaft.inertia.inertia:      0.1

    # Planetary output
    front_left_hub_shaft.inertia.inertia:       0.1
    front_right_hub_shaft.inertia.inertia:      0.1
    mid_left_hub_shaft.inertia.inertia:         0.1
    mid_right_hub_shaft.inertia.inertia:        0.1
    rear_left_hub_shaft.inertia.inertia:        0.1
    rear_right_hub_shaft.inertia.inertia:       0.1

    # Brake shafts
    front_left_brake_shaft.inertia.inertia:     0.1
    front_right_brake_shaft.inertia.inertia:    0.1
    mid_left_brake_shaft.inertia.inertia:       0.1
    mid_right_brake_shaft.inertia.inertia:      0.1
    rear_left_brake_shaft.inertia.inertia:      0.1
    rear_right_brake_shaft.inertia.inertia:     0.1


    engine becomes DriveTrain.MeanValueEngine with DriveTrain.Traits.FourStrokeEngine:
        .doc: """
        Engine data based on specifications from [1] and [2].
        """
        initial_throttle: 0.01
        torque_output.enabled: true
        rpm_output.enabled: true


        # Engine specification        
        displacement_volume:  15.24e-3  # m^3 (15.24 L)
        max_torque:     2257.0          # N·m            
        max_torque_RPM: 1400.0          # rpm            
        max_power_RPM:  2000.0          # rpm            
        idle_RPM:        700.0          # rpm            
        max_RPM:        3000.0          # rpm          

        crank_shaft_inertia:    8.0     # kg·m^2         estimated starting point for 15 L inline-6
        air_fuel_ratio:        14.5     # diesel stoich
        heat_value:      43400000.0     # J/kg diesel LHV

        # # These values are currently kept at the default settings for the MeanValueEngine.
        # idle_throttle_angle: 0.17
        # max_throttle_angle: 1.45
        # inlet_volume: engine.displacement_volume * 1.1 # 1.1× is a sensible default. 
        #                                         # If a bit more “turbo/airpath inertia” (slower boost / 
        #                                         # slower torque response), increase to 1.5×–2.0×. 
        #                                         # If response feels too sluggish, reduce to 0.7–1.0×.
        # 
        # throttle_pin_bore_ratio: 0.1    # Default. Change only if the engine feels too jumpy at small
        #                                 # throttle inputs or hard to get stable idle torque.
        # 
        # max_volumetric_efficiency: 0.92

    gear_box:
        .doc: """
        Gear box Forward ratios calculated from travel speed in each gear taken from [3] 
        
        Travel speed km/h
            Forward gears
            1st      6.6
            2nd     10.1
            3rd     15.5
            4th	    23.7
            5th	    36.5
            6th	    55.9

            Reverse gears
            1st	     6.8
            2nd	    16.0

        also using 
            tire size: 29.5R25 --> Outer diameter = 1.87 m
            differential ratio: 3.727
            final drive ratio:    4.941
            rated engine speed: 2000 rpm
        """
        initial_gear: 0
        forward_gears: [5.80, 3.79, 2.47, 1.62, 1.05, 0.69]
        reverse_gears: [5.63, 2.39]



    front_differential.gear_ratio:  3.727
    mid_differential.gear_ratio:    3.727
    rear_differential.gear_ratio:   3.727

    front_left_final_drive.ratio:   4.941
    front_right_final_drive.ratio:  4.941
    mid_left_final_drive.ratio:     4.941
    mid_right_final_drive.ratio:    4.941
    rear_left_final_drive.ratio:    4.941
    rear_right_final_drive.ratio:   4.941

    front_left_brake:
        initial_engagement_fraction: 0.0
        torque_capacity: 60e3
        min_relative_slip_ratio: 0.0001        
    front_right_brake:
        initial_engagement_fraction: 0.0
        torque_capacity: 60e3
        min_relative_slip_ratio: 0.0001        
    mid_left_brake:
        initial_engagement_fraction: 0.0
        torque_capacity: 60e3
        min_relative_slip_ratio: 0.0001        
    mid_right_brake:
        initial_engagement_fraction: 0.0
        torque_capacity: 60e3
        min_relative_slip_ratio: 0.0001        
    rear_left_brake:
        initial_engagement_fraction: 0.0
        torque_capacity: 60e3
        min_relative_slip_ratio: 0.0001        
    rear_right_brake:
        initial_engagement_fraction: 0.0
        torque_capacity: 60e3
        min_relative_slip_ratio: 0.0001        


    torque_converter:        
        .doc: """
        Torque converter is primarily based on 
        Eriksson & Nielsen - Modeling and Control of Engines and Drivelines, Section 14.8
        SAE 2004-01-0877 (heavy-duty torque converter modeling)
        Stall ratio of TC estimated from published data from other heavy machinery manufacturers
        such as Allison and ZF.    
        """

        lock_up_time: 0.3
        diameter: 0.145


SixWheelTemplate is Physics3D.System:
    # Defining all drivetrain shafts

    # Main drivetrain shafts
    engine_shaft is DriveTrain.Shaft
    torque_converter_shaft is DriveTrain.Shaft
    gear_shaft is DriveTrain.Shaft
    midrear_split_shaft is DriveTrain.Shaft

    # Differential input shafts
    front_drive_shaft is DriveTrain.Shaft
    mid_drive_shaft is DriveTrain.Shaft
    rear_drive_shaft is DriveTrain.Shaft

    # Differential output shafts
    front_left_half_shaft is DriveTrain.Shaft
    front_right_half_shaft is DriveTrain.Shaft
    mid_left_half_shaft is DriveTrain.Shaft
    mid_right_half_shaft is DriveTrain.Shaft
    rear_left_half_shaft is DriveTrain.Shaft
    rear_right_half_shaft is DriveTrain.Shaft

    # final_drive output shafts
    front_left_hub_shaft is DriveTrain.Shaft
    front_right_hub_shaft is DriveTrain.Shaft
    mid_left_hub_shaft is DriveTrain.Shaft
    mid_right_hub_shaft is DriveTrain.Shaft
    rear_left_hub_shaft is DriveTrain.Shaft
    rear_right_hub_shaft is DriveTrain.Shaft

    # Brake shafts
    front_left_brake_shaft is DriveTrain.Shaft
    front_right_brake_shaft is DriveTrain.Shaft
    mid_left_brake_shaft is DriveTrain.Shaft
    mid_right_brake_shaft is DriveTrain.Shaft
    rear_left_brake_shaft is DriveTrain.Shaft
    rear_right_brake_shaft is DriveTrain.Shaft

    # Connecting drivetrain components
    engine is DriveTrain.EmpiricalEngine or DriveTrain.MeanValueEngine:
        connectors: [engine_shaft.input]
    torque_converter is DriveTrain.EmpiricalTorqueConverter:
        connectors: [engine_shaft.output, torque_converter_shaft.input]
    gear_box is DriveTrain.GearBox:
        connectors: [torque_converter_shaft.output, gear_shaft.input]

    # Using double ordinary differentials to connect to both mid axle and to the rear differential.
    mid_center_axle is DriveTrain.Differential:
        drive_shaft: gear_shaft
        left_axle_shaft: front_drive_shaft
        right_axle_shaft: midrear_split_shaft
        locked_up: true
        locked_up_output.enabled: true
    rear_center_axle is DriveTrain.Differential:
        drive_shaft: midrear_split_shaft
        left_axle_shaft: mid_drive_shaft
        right_axle_shaft: rear_drive_shaft
        locked_up: true
        locked_up_output.enabled: true

    # Axle differentials
    front_differential is DriveTrain.Differential:
        drive_shaft: front_drive_shaft
        left_axle_shaft: front_left_half_shaft
        right_axle_shaft: front_right_half_shaft
        locked_up: false
        locked_up_output.enabled: true
    mid_differential is DriveTrain.Differential:
        drive_shaft: mid_drive_shaft
        left_axle_shaft: mid_left_half_shaft
        right_axle_shaft: mid_right_half_shaft
        locked_up: false
        locked_up_output.enabled: true
    rear_differential is DriveTrain.Differential:
        drive_shaft: rear_drive_shaft
        left_axle_shaft: rear_left_half_shaft
        right_axle_shaft: rear_right_half_shaft
        locked_up: false
        locked_up_output.enabled: true

    # Brakes
    front_left_brake is DriveTrain.ManualBrake:
        connectors: [front_left_half_shaft.output, front_left_brake_shaft.input]
    front_right_brake is DriveTrain.ManualBrake:
        connectors: [front_right_half_shaft.output, front_right_brake_shaft.input]
    mid_left_brake is DriveTrain.ManualBrake:
        connectors: [mid_left_half_shaft.output, mid_left_brake_shaft.input]
    mid_right_brake is DriveTrain.ManualBrake:
        connectors: [mid_right_half_shaft.output, mid_right_brake_shaft.input]
    rear_left_brake is DriveTrain.ManualBrake:
        connectors: [rear_left_half_shaft.output, rear_left_brake_shaft.input]
    rear_right_brake is DriveTrain.ManualBrake:
        connectors: [rear_right_half_shaft.output, rear_right_brake_shaft.input]

    # Hub final drive reductions
    front_left_final_drive is DriveTrain.Gear:
        connectors: [front_left_brake_shaft.output, front_left_hub_shaft.input]
    front_right_final_drive is DriveTrain.Gear:
        connectors: [front_right_brake_shaft.output, front_right_hub_shaft.input]
    mid_left_final_drive is DriveTrain.Gear:
        connectors: [mid_left_brake_shaft.output, mid_left_hub_shaft.input]
    mid_right_final_drive is DriveTrain.Gear:
        connectors: [mid_right_brake_shaft.output, mid_right_hub_shaft.input]
    rear_left_final_drive is DriveTrain.Gear:
        connectors: [rear_left_brake_shaft.output, rear_left_hub_shaft.input]
    rear_right_final_drive is DriveTrain.Gear:
        connectors: [rear_right_brake_shaft.output, rear_right_hub_shaft.input]

    # Wheel actuators
    front_left_actuator is DriveTrain.Actuator:
        connector_1d: front_left_hub_shaft.output
    front_right_actuator is DriveTrain.Actuator:
        connector_1d: front_right_hub_shaft.output
    mid_left_actuator is DriveTrain.Actuator:
        connector_1d: mid_left_hub_shaft.output
    mid_right_actuator is DriveTrain.Actuator:
        connector_1d: mid_right_hub_shaft.output
    rear_left_actuator is DriveTrain.Actuator:
        connector_1d: rear_left_hub_shaft.output
    rear_right_actuator is DriveTrain.Actuator:
        connector_1d: rear_right_hub_shaft.output

