#!/usr/bin/env python3
"""
Gurobipy 12.0.2 Implementation for Music Storage Optimization
"""

import gurobipy as gp
from gurobipy import GRB

def music_storage_optimization():
    """Optimize music storage to minimize total file size while meeting constraints."""
    
    # 1. MODEL & DATA SETUP
    model = gp.Model("music_storage_optimization")
    
    # Example data (replace with actual data from database)
    songs = [
        {'id': 1, 'file_size': 8, 'rating': 4, 'artist': 'ArtistA', 'genre': 'Pop'},
        {'id': 2, 'file_size': 12, 'rating': 5, 'artist': 'ArtistB', 'genre': 'Rock'},
        {'id': 3, 'file_size': 6, 'rating': 3, 'artist': 'ArtistA', 'genre': 'Pop'}
    ]
    
    constraints = {
        'min_total_songs': 2,
        'min_avg_rating': 4,
        'max_songs_per_artist': 1,
        'min_songs_per_genre': 1
    }
    
    # Extract unique artists and genres
    artists = list(set(song['artist'] for song in songs))
    genres = list(set(song['genre'] for song in songs))
    
    # CRITICAL: Validate array lengths before loops
    assert len(songs) > 0, "No songs provided"
    assert len(artists) > 0, "No artists provided"
    assert len(genres) > 0, "No genres provided"
    
    # 2. VARIABLES
    x = {song['id']: model.addVar(vtype=GRB.BINARY, name=f"x_{song['id']}") for song in songs}
    
    # 3. OBJECTIVE FUNCTION
    model.setObjective(gp.quicksum(song['file_size'] * x[song['id']] for song in songs), GRB.MINIMIZE)
    
    # 4. CONSTRAINTS
    
    # Minimum Total Songs Stored
    model.addConstr(gp.quicksum(x[song['id']] for song in songs) >= constraints['min_total_songs'], name="min_total_songs")
    
    # Minimum Average Rating (Linearized)
    model.addConstr(gp.quicksum((song['rating'] - constraints['min_avg_rating']) * x[song['id']] for song in songs) >= 0, name="min_avg_rating")
    
    # Maximum Songs per Artist
    for artist in artists:
        artist_songs = [song['id'] for song in songs if song['artist'] == artist]
        model.addConstr(gp.quicksum(x[song_id] for song_id in artist_songs) <= constraints['max_songs_per_artist'], name=f"max_songs_artist_{artist}")
    
    # Minimum Songs per Genre
    for genre in genres:
        genre_songs = [song['id'] for song in songs if song['genre'] == genre]
        model.addConstr(gp.quicksum(x[song_id] for song_id in genre_songs) >= constraints['min_songs_per_genre'], name=f"min_songs_genre_{genre}")
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for song in songs:
            if x[song['id']].x > 1e-6:
                print(f"Song {song['id']} is stored locally.")
    elif model.status == GRB.INFEASIBLE:
        print("Problem is infeasible")
    elif model.status == GRB.UNBOUNDED:
        print("Problem is unbounded")
    
    return model

# Run the optimization
if __name__ == "__main__":
    music_storage_optimization()