#!/usr/bin/env python3
"""
DOCPLEX implementation for music storage optimization problem
"""

from docplex.mp.model import Model

def music_storage_optimization():
    """Optimize music storage to minimize total file size while meeting constraints"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="music_storage_optimization")
    
    # Example data (replace with actual data loading)
    file_sizes = [8, 12, 6]  # File sizes in MB
    ratings = [4, 5, 3]      # Song ratings
    artists = [1, 2, 1]      # Artist IDs
    genres = [1, 2, 1]       # Genre IDs
    
    # Constraints
    min_total_songs = 2
    min_avg_rating = 4
    max_songs_per_artist = 1
    min_songs_per_genre = 1
    
    # CRITICAL: Validate array lengths to prevent IndexError
    assert len(file_sizes) == len(ratings) == len(artists) == len(genres), "Array length mismatch"
    safe_range = range(min(len(file_sizes), len(ratings), len(artists), len(genres)))  # Safe indexing
    
    # 2. VARIABLES
    x = {i: mdl.binary_var(name=f"x_{i}") for i in safe_range}  # Decision variables
    
    # 3. OBJECTIVE FUNCTION
    total_file_size = mdl.sum(file_sizes[i] * x[i] for i in safe_range)
    mdl.minimize(total_file_size)
    
    # 4. CONSTRAINTS
    
    # Minimum Total Songs Stored
    total_songs = mdl.sum(x[i] for i in safe_range)
    mdl.add_constraint(total_songs >= min_total_songs, ctname="min_total_songs")
    
    # Minimum Average Rating (linearized)
    total_rating = mdl.sum((ratings[i] - min_avg_rating) * x[i] for i in safe_range)
    mdl.add_constraint(total_rating >= 0, ctname="min_avg_rating")
    
    # Maximum Songs per Artist
    for artist in set(artists):
        artist_songs = mdl.sum(x[i] for i in safe_range if artists[i] == artist)
        mdl.add_constraint(artist_songs <= max_songs_per_artist, ctname=f"max_songs_artist_{artist}")
    
    # Minimum Songs per Genre
    for genre in set(genres):
        genre_songs = mdl.sum(x[i] for i in safe_range if genres[i] == genre)
        mdl.add_constraint(genre_songs >= min_songs_per_genre, ctname=f"min_songs_genre_{genre}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for i in safe_range:
            value = solution.get_value(x[i])
            if value > 1e-6:
                print(f"Song {i} is stored locally: {value}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

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