# -- DATA SECTION: start

@kawa_tool(
    outputs={
        'trade_id': str,
        'trade_date': date,
        'commodity': str,
        'instrument_type': str,
        'desk': str,
        'region': str,
        'counterparty': str,
        'trader_id': str,
        'quantity': float,
        'price': float,
        'currency': str,
        'delivery_location': str,
        'delivery_month': str
    }
)
def trades_generator():
    from faker import Faker
    import pandas as pd
    import numpy as np
    from datetime import date, timedelta
    
    fake = Faker()
    
    commodities = ['Crude_WTI', 'Crude_Brent', 'Gas_HenryHub', 'Gas_TTF', 'LNG_JKM', 'Gasoline_RBOB', 'Diesel_Gasoil']
    instrument_types = ['Physical', 'Futures', 'Swap', 'Option']
    instrument_weights = [0.40, 0.35, 0.20, 0.05]
    
    desks = ['Houston_Crude', 'London_Crude', 'Singapore_Crude', 'Houston_Gas', 'London_Gas', 'Tokyo_LNG', 'NewYork_Refining']
    counterparties = ['Shell', 'BP', 'Vitol', 'Trafigura', 'TotalEnergies', 'Chevron', 'ExxonMobil', 'Glencore', 'Gunvor', 'Koch']
    
    desk_to_region = {
        'Houston_Crude': 'Americas',
        'London_Crude': 'Europe',
        'Singapore_Crude': 'Asia',
        'Houston_Gas': 'Americas',
        'London_Gas': 'Europe',
        'Tokyo_LNG': 'Asia',
        'NewYork_Refining': 'Americas'
    }
    
    desk_to_commodities = {
        'Houston_Crude': ['Crude_WTI'],
        'London_Crude': ['Crude_Brent', 'Diesel_Gasoil'],
        'Singapore_Crude': ['Crude_Brent'],
        'Houston_Gas': ['Gas_HenryHub'],
        'London_Gas': ['Gas_TTF'],
        'Tokyo_LNG': ['LNG_JKM'],
        'NewYork_Refining': ['Gasoline_RBOB', 'Diesel_Gasoil']
    }
    
    commodity_to_location = {
        'Crude_WTI': ['Cushing_OK'],
        'Crude_Brent': ['Rotterdam', 'Singapore'],
        'Gas_HenryHub': ['HenryHub_LA'],
        'Gas_TTF': ['Rotterdam'],
        'LNG_JKM': ['Tokyo_Bay', 'Singapore'],
        'Gasoline_RBOB': ['NewYork_Harbor'],
        'Diesel_Gasoil': ['Rotterdam', 'Singapore']
    }
    
    price_ranges = {
        'Crude_WTI': (70.0, 72.5),
        'Crude_Brent': (76.0, 78.0),
        'Gas_HenryHub': (2.70, 2.90),
        'Gas_TTF': (9.0, 9.6),
        'LNG_JKM': (13.3, 14.2),
        'Gasoline_RBOB': (2.00, 2.15),
        'Diesel_Gasoil': (710.0, 730.0)
    }
    
    quantity_ranges = {
        'Crude_WTI': (250000, 1000000),
        'Crude_Brent': (250000, 1000000),
        'Gas_HenryHub': (8000000, 15000000),
        'Gas_TTF': (10000000, 16000000),
        'LNG_JKM': (45000, 75000),
        'Gasoline_RBOB': (80000, 200000),
        'Diesel_Gasoil': (150000, 250000)
    }
    
    data = []
    start_date = date(2024, 1, 2)
    trader_ids = [f'TRD_{i:03d}' for i in range(1, 51)]
    
    for i in range(300):
        trade_id = f'TRD_{i+1:05d}'
        
        days_offset = np.random.randint(0, 25)
        trade_date = start_date + timedelta(days=days_offset)
        while trade_date.weekday() >= 5:
            trade_date = trade_date + timedelta(days=1)
        
        desk = np.random.choice(desks)
        region = desk_to_region[desk]
        commodity = np.random.choice(desk_to_commodities[desk])
        instrument_type = np.random.choice(instrument_types, p=instrument_weights)
        
        counterparty = np.random.choice(counterparties)
        trader_id = np.random.choice(trader_ids)
        
        price_min, price_max = price_ranges[commodity]
        price = np.random.uniform(price_min, price_max)
        price = round(price + np.random.normal(0, (price_max - price_min) * 0.03), 2)
        
        qty_min, qty_max = quantity_ranges[commodity]
        quantity = np.random.uniform(qty_min, qty_max)
        if np.random.random() < 0.15:
            quantity *= np.random.uniform(1.3, 2.2)
        quantity = round(quantity, 0)
        
        currency = 'USD'
        if desk == 'London_Gas':
            currency = 'EUR'
        elif np.random.random() < 0.08:
            currency = np.random.choice(['GBP', 'JPY', 'SGD'])
        
        delivery_location = np.random.choice(commodity_to_location[commodity])
        
        delivery_month_offset = np.random.choice([1, 2, 3], p=[0.6, 0.3, 0.1])
        delivery_month_date = trade_date.replace(day=1)
        for _ in range(delivery_month_offset):
            if delivery_month_date.month == 12:
                delivery_month_date = delivery_month_date.replace(year=delivery_month_date.year + 1, month=1)
            else:
                delivery_month_date = delivery_month_date.replace(month=delivery_month_date.month + 1)
        delivery_month = delivery_month_date.strftime('%Y-%m')
        
        data.append([
            trade_id, trade_date, commodity, instrument_type, desk, region,
            counterparty, trader_id, quantity, price, currency,
            delivery_location, delivery_month
        ])
    
    df = pd.DataFrame(data, columns=[
        'trade_id', 'trade_date', 'commodity', 'instrument_type', 'desk', 'region',
        'counterparty', 'trader_id', 'quantity', 'price', 'currency',
        'delivery_location', 'delivery_month'
    ])
    
    return df


trades_dataset = app.create_dataset(
    name='Trades',
    generator=trades_generator,
)


@kawa_tool(
    outputs={
        'date': date,
        'commodity': str,
        'region': str,
        'benchmark': str,
        'price': float,
        'volatility': float
    }
)
def market_prices_generator():
    import pandas as pd
    import numpy as np
    from datetime import date, timedelta
    
    commodities_config = [
        ('Crude_WTI', 'Americas', 'NYMEX_CL_Front', 71.2, 0.45, 22.8),
        ('Crude_Brent', 'Europe', 'ICE_Brent_Front', 77.0, 0.48, 21.9),
        ('Crude_Brent', 'Asia', 'ICE_Brent_Front', 77.2, 0.50, 22.1),
        ('Gas_HenryHub', 'Americas', 'NYMEX_NG_Front', 2.81, 0.048, 35.8),
        ('Gas_TTF', 'Europe', 'ICE_TTF_Front', 9.35, 0.16, 41.8),
        ('LNG_JKM', 'Asia', 'Platts_JKM', 13.72, 0.22, 28.9),
        ('Gasoline_RBOB', 'Americas', 'NYMEX_RBOB_Front', 2.07, 0.032, 26.3),
        ('Diesel_Gasoil', 'Europe', 'ICE_Gasoil_Front', 719.5, 6.2, 24.5),
        ('Diesel_Gasoil', 'Asia', 'Platts_Gasoil', 721.0, 6.5, 24.8)
    ]
    
    data = []
    start_date = date(2024, 1, 2)
    num_days = 25
    
    for commodity, region, benchmark, base_price, daily_vol, base_volatility in commodities_config:
        price = base_price
        volatility = base_volatility
        
        for day_offset in range(num_days):
            current_date = start_date + timedelta(days=day_offset)
            
            if current_date.weekday() < 5:
                price_change = np.random.normal(0, daily_vol)
                trend_factor = np.sin(day_offset / 5.0) * daily_vol * 0.3
                price = price + price_change + trend_factor
                
                volatility = base_volatility + np.random.uniform(-2.5, 2.5)
                volatility = max(15.0, min(50.0, volatility))
                
                if np.random.random() < 0.05:
                    price += np.random.choice([-1, 1]) * daily_vol * 3
                
                data.append([
                    current_date,
                    commodity,
                    region,
                    benchmark,
                    round(price, 2),
                    round(volatility, 1)
                ])
    
    df = pd.DataFrame(data, columns=[
        'date', 'commodity', 'region', 'benchmark', 'price', 'volatility'
    ])
    
    return df


market_prices_dataset = app.create_dataset(
    name='Market Prices',
    generator=market_prices_generator,
)


@kawa_tool(
    outputs={
        'trade_id': str,
        'date': date,
        'daily_pnl': float,
        'price_effect': float,
        'volume_effect': float,
        'fx_effect': float,
        'basis_effect': float,
        'unexplained_pnl': float
    }
)
def pnl_attribution_generator():
    import pandas as pd
    import numpy as np
    from datetime import date, timedelta
    
    np.random.seed(42)
    
    sample_trades = [
        ('TRD_00001', date(2024, 1, 2), 'Crude_WTI', 500000, 71.25, 'Physical'),
        ('TRD_00002', date(2024, 1, 2), 'Crude_Brent', 1000000, 76.80, 'Futures'),
        ('TRD_00003', date(2024, 1, 3), 'Gas_HenryHub', 10000000, 2.85, 'Swap'),
        ('TRD_00004', date(2024, 1, 3), 'LNG_JKM', 70000, 13.50, 'Physical'),
        ('TRD_00005', date(2024, 1, 4), 'Gas_TTF', 15000000, 9.20, 'Futures'),
        ('TRD_00006', date(2024, 1, 4), 'Crude_WTI', 250000, 72.00, 'Option'),
        ('TRD_00007', date(2024, 1, 5), 'Gasoline_RBOB', 100000, 2.05, 'Physical'),
        ('TRD_00008', date(2024, 1, 5), 'Crude_Brent', 750000, 77.50, 'Physical'),
        ('TRD_00009', date(2024, 1, 8), 'Diesel_Gasoil', 200000, 720.00, 'Futures'),
        ('TRD_00010', date(2024, 1, 8), 'LNG_JKM', 50000, 13.75, 'Swap'),
        ('TRD_00011', date(2024, 1, 9), 'Crude_WTI', 800000, 70.90, 'Futures'),
        ('TRD_00012', date(2024, 1, 9), 'Gas_TTF', 12000000, 9.45, 'Physical'),
        ('TRD_00013', date(2024, 1, 10), 'Crude_Brent', 600000, 76.20, 'Swap'),
        ('TRD_00014', date(2024, 1, 10), 'Gas_HenryHub', 8000000, 2.78, 'Physical'),
        ('TRD_00015', date(2024, 1, 11), 'LNG_JKM', 65000, 14.00, 'Physical'),
        ('TRD_00016', date(2024, 1, 11), 'Gasoline_RBOB', 150000, 2.08, 'Futures'),
        ('TRD_00017', date(2024, 1, 12), 'Crude_WTI', 550000, 71.80, 'Physical'),
        ('TRD_00018', date(2024, 1, 12), 'Diesel_Gasoil', 180000, 715.00, 'Physical'),
        ('TRD_00019', date(2024, 1, 15), 'Crude_Brent', 900000, 77.10, 'Futures'),
        ('TRD_00020', date(2024, 1, 15), 'Gas_TTF', 11000000, 9.30, 'Swap'),
    ]
    
    data = []
    
    for trade_id, trade_date, commodity, quantity, entry_price, instrument_type in sample_trades:
        
        if instrument_type == 'Option':
            num_days = np.random.randint(8, 15)
        elif instrument_type == 'Physical':
            num_days = np.random.randint(12, 20)
        else:
            num_days = np.random.randint(10, 18)
        
        cumulative_loss_factor = 0.0
        if trade_id in ['TRD_00008', 'TRD_00006']:
            cumulative_loss_factor = -0.02
        
        for day_offset in range(1, num_days + 1):
            pnl_date = trade_date + timedelta(days=day_offset)
            
            if pnl_date.weekday() >= 5:
                continue
            
            base_volatility = 0.018 if commodity.startswith('Crude') else 0.025
            price_change_pct = np.random.normal(cumulative_loss_factor, base_volatility)
            
            if commodity.startswith('Crude'):
                price_effect = price_change_pct * entry_price * quantity
            elif commodity.startswith('Gas'):
                price_effect = price_change_pct * entry_price * quantity * 0.0012
            elif commodity.startswith('LNG'):
                price_effect = price_change_pct * entry_price * quantity * 0.18
            elif commodity.startswith('Gasoline'):
                price_effect = price_change_pct * entry_price * quantity * 0.85
            else:
                price_effect = price_change_pct * entry_price * quantity * 0.75
            
            volume_effect = 0.0
            if np.random.random() < 0.08:
                volume_effect = np.random.uniform(-0.03, 0.03) * abs(price_effect)
            
            fx_effect = 0.0
            if np.random.random() < 0.18:
                fx_effect = np.random.uniform(-0.10, 0.10) * abs(price_effect)
            
            basis_volatility = 0.012 if instrument_type == 'Physical' else 0.008
            basis_effect = np.random.uniform(-basis_volatility, basis_volatility) * abs(price_effect)
            
            if instrument_type == 'Option':
                theta_decay = -abs(price_effect) * 0.05 * (day_offset / num_days)
                price_effect += theta_decay
            
            daily_pnl = price_effect + volume_effect + fx_effect + basis_effect
            
            unexplained_pct = np.random.uniform(0.008, 0.028)
            unexplained_pnl = np.random.choice([-1, 1]) * unexplained_pct * abs(daily_pnl)
            daily_pnl += unexplained_pnl
            
            data.append([
                trade_id,
                pnl_date,
                round(daily_pnl, 2),
                round(price_effect, 2),
                round(volume_effect, 2),
                round(fx_effect, 2),
                round(basis_effect, 2),
                round(unexplained_pnl, 2)
            ])
    
    df = pd.DataFrame(data, columns=[
        'trade_id', 'date', 'daily_pnl', 'price_effect', 'volume_effect',
        'fx_effect', 'basis_effect', 'unexplained_pnl'
    ])
    
    return df


pnl_attribution_dataset = app.create_dataset(
    name='PnL Attribution',
    generator=pnl_attribution_generator,
)


@kawa_tool(
    outputs={
        'desk': str,
        'date': date,
        'var_amount': float,
        'stress_pnl': float,
        'exposure': float,
        'limit_amount': float,
        'limit_breach_flag': float
    }
)
def risk_metrics_generator():
    import pandas as pd
    import numpy as np
    from datetime import date, timedelta
    
    desks = ['Houston_Crude', 'London_Crude', 'Singapore_Crude', 'Houston_Gas', 'London_Gas', 'Tokyo_LNG', 'NewYork_Refining']
    
    desk_config = {
        'Houston_Crude': {'base_var': 1300000, 'limit': 1500000, 'base_exposure': 88.5, 'volatility': 0.06},
        'London_Crude': {'base_var': 1875000, 'limit': 2000000, 'base_exposure': 130.2, 'volatility': 0.055},
        'Singapore_Crude': {'base_var': 1025000, 'limit': 1200000, 'base_exposure': 68.8, 'volatility': 0.065},
        'Houston_Gas': {'base_var': 715000, 'limit': 800000, 'base_exposure': 49.5, 'volatility': 0.075},
        'London_Gas': {'base_var': 1165000, 'limit': 1200000, 'base_exposure': 77.8, 'volatility': 0.080},
        'Tokyo_LNG': {'base_var': 835000, 'limit': 1000000, 'base_exposure': 57.2, 'volatility': 0.068},
        'NewYork_Refining': {'base_var': 445000, 'limit': 500000, 'base_exposure': 31.2, 'volatility': 0.072}
    }
    
    data = []
    start_date = date(2024, 1, 2)
    num_days = 25
    
    breach_schedule = {
        ('London_Gas', 7): 1.05,
        ('London_Gas', 8): 1.08,
        ('London_Crude', 9): 1.06,
        ('Houston_Gas', 11): 1.04
    }
    
    for desk in desks:
        config = desk_config[desk]
        base_var = config['base_var']
        limit_amount = config['limit']
        base_exposure = config['base_exposure']
        var_volatility = config['volatility']
        
        var_amount = base_var
        exposure = base_exposure
        
        for day_offset in range(num_days):
            current_date = start_date + timedelta(days=day_offset)
            
            if current_date.weekday() >= 5:
                continue
            
            if (desk, day_offset) in breach_schedule:
                var_amount = limit_amount * breach_schedule[(desk, day_offset)]
            else:
                var_change = np.random.normal(0, base_var * var_volatility)
                trend = np.sin(day_offset / 4.0) * base_var * 0.03
                var_amount = var_amount + var_change + trend
                
                var_amount = max(base_var * 0.55, min(base_var * 1.45, var_amount))
                
                if np.random.random() < 0.15 and day_offset > 5:
                    var_amount = limit_amount * np.random.uniform(0.82, 0.95)
            
            exposure_change = np.random.normal(0, base_exposure * 0.045)
            exposure = exposure + exposure_change
            exposure = max(base_exposure * 0.65, min(base_exposure * 1.38, exposure))
            
            stress_multiplier = np.random.uniform(2.15, 2.55)
            stress_pnl = -var_amount * stress_multiplier
            
            limit_breach_flag = 1.0 if var_amount >= limit_amount else 0.0
            
            data.append([
                desk,
                current_date,
                round(var_amount, 0),
                round(stress_pnl, 0),
                round(exposure, 1),
                round(limit_amount, 0),
                limit_breach_flag
            ])
    
    df = pd.DataFrame(data, columns=[
        'desk', 'date', 'var_amount', 'stress_pnl', 'exposure', 'limit_amount', 'limit_breach_flag'
    ])
    
    return df


risk_metrics_dataset = app.create_dataset(
    name='Risk Metrics',
    generator=risk_metrics_generator,
)

model = app.create_model(
    dataset=trades_dataset,
)

# -- DATA SECTION: end

# -- RELATIONSHIPS SECTION: start

pnl_relationship = model.create_relationship(
    name='PnL Attribution',
    dataset=pnl_attribution_dataset,
    link={'trade_id': 'trade_id'}
)

pnl_relationship.add_column(
    name='daily_pnl',
    aggregation='SUM',
    new_column_name='total_pnl',
)

pnl_relationship.add_column(
    name='price_effect',
    aggregation='SUM',
    new_column_name='total_price_effect',
)

pnl_relationship.add_column(
    name='volume_effect',
    aggregation='SUM',
    new_column_name='total_volume_effect',
)

pnl_relationship.add_column(
    name='fx_effect',
    aggregation='SUM',
    new_column_name='total_fx_effect',
)

pnl_relationship.add_column(
    name='basis_effect',
    aggregation='SUM',
    new_column_name='total_basis_effect',
)

pnl_relationship.add_column(
    name='unexplained_pnl',
    aggregation='SUM',
    new_column_name='total_unexplained_pnl',
)

pnl_relationship.add_column(
    name='daily_pnl',
    aggregation='COUNT',
    new_column_name='pnl_days_count',
)

market_prices_relationship = model.create_relationship(
    name='Market Prices',
    dataset=market_prices_dataset,
    link={'commodity': 'commodity', 'region': 'region'}
)

market_prices_relationship.add_column(
    name='price',
    aggregation='AVERAGE',
    new_column_name='avg_market_price',
)

market_prices_relationship.add_column(
    name='volatility',
    aggregation='AVERAGE',
    new_column_name='avg_market_volatility',
)

market_prices_relationship.add_column(
    name='price',
    aggregation='MIN',
    new_column_name='min_market_price',
)

market_prices_relationship.add_column(
    name='price',
    aggregation='MAX',
    new_column_name='max_market_price',
)

risk_metrics_relationship = model.create_relationship(
    name='Risk Metrics',
    dataset=risk_metrics_dataset,
    link={'desk': 'desk'}
)

risk_metrics_relationship.add_column(
    name='var_amount',
    aggregation='AVERAGE',
    new_column_name='avg_var',
)

risk_metrics_relationship.add_column(
    name='stress_pnl',
    aggregation='AVERAGE',
    new_column_name='avg_stress_pnl',
)

risk_metrics_relationship.add_column(
    name='exposure',
    aggregation='AVERAGE',
    new_column_name='avg_exposure',
)

risk_metrics_relationship.add_column(
    name='limit_amount',
    aggregation='FIRST',
    new_column_name='desk_limit',
)

risk_metrics_relationship.add_column(
    name='limit_breach_flag',
    aggregation='SUM',
    new_column_name='total_breach_days',
)

risk_metrics_relationship.add_column(
    name='var_amount',
    aggregation='MAX',
    new_column_name='max_var',
)

# -- RELATIONSHIPS SECTION: end
