import streamlit as st
import pandas as pd

st.title("Contract Balance Allocator")

# Configuration section
st.header("Configuration")

# Number of contracts
num_contracts = st.number_input(
    "Number of Contracts",
    min_value=2,
    max_value=20,
    value=5,
    step=1
)

# Initialize session state for contract balances if not exists
if 'balances' not in st.session_state or len(st.session_state.balances) != num_contracts:
    st.session_state.balances = [0.0] * num_contracts

# Input balances for each contract
st.subheader("Contract Balances")
cols = st.columns(min(3, num_contracts))
for i in range(num_contracts):
    col_idx = i % 3
    with cols[col_idx]:
        st.session_state.balances[i] = st.number_input(
            f"Contract {i+1}",
            value=st.session_state.balances[i],
            step=100.0,
            format="%.2f",
            key=f"contract_{i}"
        )

# Display current state
st.header("Current State")
current_df = pd.DataFrame({
    'Contract': [f"Contract {i+1}" for i in range(num_contracts)],
    'Balance': st.session_state.balances
})
st.dataframe(current_df, hide_index=True)

# Allocation logic
def allocate_balances(balances):
    """
    Allocate positive balances to negative balances.
    Uses largest positive balance first, then next largest, etc.
    For negative balances, prioritize those closest to 0 first.
    """
    # Create a working copy with contract indices
    contracts = [(i, balance) for i, balance in enumerate(balances)]

    # Separate positive and negative balances
    positive_contracts = [(i, bal) for i, bal in contracts if bal > 0]

    # Sort positive balances in descending order (largest first)
    positive_contracts.sort(key=lambda x: x[1], reverse=True)

    # Create result array
    result = balances.copy()

    # Track allocations for display
    allocations = []

    # Process each positive balance
    for pos_idx, pos_balance in positive_contracts:
        remaining = pos_balance

        # Keep allocating until the positive balance is exhausted
        while remaining > 0:
            # Find all contracts with negative balances
            eligible_negatives = [(i, result[i]) for i in range(len(result)) if result[i] < 0]

            # If no negative balances remain, stop
            if not eligible_negatives:
                break

            # Sort by balance descending (closest to 0 first)
            # e.g., -100 comes before -500
            eligible_negatives.sort(key=lambda x: x[1], reverse=True)

            # Take the negative balance closest to 0
            neg_idx, neg_balance = eligible_negatives[0]

            # Calculate how much we can allocate (don't exceed 0)
            allocation = min(remaining, abs(neg_balance))

            # Apply allocation
            result[pos_idx] -= allocation
            result[neg_idx] += allocation
            remaining -= allocation

            # Record allocation
            allocations.append({
                'From': f"Contract {pos_idx+1}",
                'To': f"Contract {neg_idx+1}",
                'Amount': allocation
            })

    return result, allocations

# Allocate button
if st.button("Allocate Balances", type="primary"):
    new_balances, allocations = allocate_balances(st.session_state.balances)

    # Display allocations
    st.header("Allocation Details")
    if allocations:
        allocation_df = pd.DataFrame(allocations)
        st.dataframe(allocation_df, hide_index=True)
    else:
        st.info("No allocations needed. No positive balances to allocate or no negative balances to receive.")

    # Display results
    st.header("Results")
    result_df = pd.DataFrame({
        'Contract': [f"Contract {i+1}" for i in range(num_contracts)],
        'Original Balance': st.session_state.balances,
        'Final Balance': new_balances,
        'Change': [new - old for new, old in zip(new_balances, st.session_state.balances)]
    })
    st.dataframe(result_df, hide_index=True)

    # Summary statistics
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("Total Original", f"${sum(st.session_state.balances):,.2f}")
    with col2:
        st.metric("Total Final", f"${sum(new_balances):,.2f}")
    with col3:
        remaining_positive = sum(b for b in new_balances if b > 0)
        st.metric("Remaining Positive", f"${remaining_positive:,.2f}")