import copy
import json
from collections import defaultdict
from .models import LocationGroupModel, alloction_pre
from .queries import LocationQueries

class AllocationAlgorithm:
    CAPACITY_MAP = {'T1':1, 'T2':2, 'T4':4, 'S4':4, 'T5':5}

    @classmethod
    def generate_plan(cls, total_pallets, layer_capacity, current_pressure):
        def allocate(remain, path, pressure, layer_state, depth=0):
            if remain <= 0:
                return path, pressure
                
            balance_factor = 1.0 - (0.1 * min(depth, 5))
            layer_priority = sorted(
                enumerate(pressure),
                key=lambda x: (x[1] * balance_factor, x[0])
            )

            for layer_idx, _ in layer_priority:
                candidates = sorted(
                    [(t, c) for t, c in cls.CAPACITY_MAP.items()
                     if layer_state[layer_idx].get(t, 0) > 0],
                    key=lambda x: (abs(x[1]-remain), -x[1])
                )

                for loc_type, cap in candidates:
                    new_state = copy.deepcopy(layer_state)
                    new_state[layer_idx][loc_type] -= 1
                    
                    # new_pressure = [p-1 if p>0 else 0 for p in pressure]
                    new_pressure = [p if p>0 else 0 for p in pressure]
                    allocated = min(cap, remain)
                    new_pressure[layer_idx] += allocated

                    result = allocate(
                        remain - allocated,
                        path + [f"{layer_idx+1}_{loc_type}"],
                        new_pressure,
                        new_state,
                        depth + 1
                    )
                    if result:
                        return result
            return None

        return allocate(total_pallets, [], current_pressure, layer_capacity)

    @staticmethod
    def format_solution(solution):
        result = defaultdict(lambda: defaultdict(int))
        for item in solution:
            parts = item.split('_')
            if len(parts) == 2:
                layer, loc_type = parts
                result[layer][loc_type] += 1
        return json.dumps(result)
    
    @staticmethod
    def allocation_plan_left_right(location_type_list,batch_number,start_location,container_code):
        locations=[]
        existing_solution = alloction_pre.objects.filter(batch_number=batch_number).first()

        for layer, location_type_dict in location_type_list.items():
            if not location_type_dict:
                continue
            # 获取库位类型列表
            location_type = list(location_type_dict.keys())
            demand_number = sum(location_type_dict.values())
            print (f"[1]层{layer} 需求数量: {demand_number}, 库位: {location_type}")                    
            location_groups = LocationGroupModel.objects.filter(
                group_type__in=location_type,
                layer=layer,
                status='available'
            )
            if not location_groups:
                print(f"层{layer} 无库位")
            # 根据起始位置选择排序字段
            if start_location == '203':
                ordered_groups = location_groups.order_by('left_priority')
            elif start_location == '103':
                ordered_groups = location_groups.order_by('right_priority')
            else:
                ordered_groups = location_groups.none()

            number = 0
            for location_group in ordered_groups:
                if number >= demand_number:
                    break
                locations.append(f"{layer}_{location_group.id}")
                number += 1
        
        existing_solution.layer_solution_type = locations
        existing_solution.save()
        print(f"[2]分配方案: {locations}")
        return locations if locations else None
      
    def generate_WCS_location(location_list_cnumber):
        allocation_target_location = (
                            location_list_cnumber.warehouse_code + '-' 
                            + f"{int(location_list_cnumber.row):02d}" + '-'  
                            + f"{int(location_list_cnumber.col):02d}" + '-' 
                            + f"{int(location_list_cnumber.layer):02d}" 
                        )
        return allocation_target_location