#!/usr/local/bin/python3
# assign.py : Assign people to teams
#
# Code by: ajpawale-cdeshpa-spanampi-a1
#
# Based on skeleton code by R. Shah and D. Crandall, January 2021
#

import sys
import time
import random
import copy
from itertools import combinations

dictionary={}
dictionary_cost={}
dictionary_cost_features={}
keys=[]
requested=[]
not_wanted=[]
indices=[]

"""Initially we convert the details of the input file into 3 arrays: keys(containing individual names), requested(containing the team members which are requested), 
not_wanted(containing names of team members they don't want to work with) and indices(containing the indices of keys to track the count of cost function. 
This has been handled using fopen() in Python.

Along with this we have also initialized some dictionaries:
a. dictionary: To store the information about the person along with his requested teammates and non-wanted teammates.
b. dictionary_cost: To keep a track of the cost generated by each person(mails).
c. dictionary_cost_features: To keep a track of the length of the team members requested, team members requested (used to evaluate cost function) and not wanted team members.
"""
def initialize(input_file):

    f=open(input_file,"r")
    for line in f.readlines():
        for i in line.split():
            keys.append(i)							#Append the 1st column from file to keys list.
            break
    f.close()

    f=open(input_file,"r") 
    for line in f.readlines():
        for j in line.split():
            requested.append(line.split()[1])		#Append the 2nd column from file to requested list
            break

    f.close()

    f=open(input_file,"r") 
    for line in f.readlines():
        for j in line.split():
            not_wanted.append(line.split()[2])		#Append the 3rd column from file to not_wanted list
            break

    f.close()

    for i in range(len(keys)):
        indices.append(i)


    zip_iterator = zip(indices, keys)
    dictionary_indices = dict(zip_iterator)    


    for i in range(len(keys)):
        dictionary[keys[i]]=[[],[]]
        dictionary_cost[keys[i]]=0
        dictionary_cost_features[keys[i]]=[0,[]]


    for key,value in dictionary.items():
        for i in range(len(keys)):
            dictionary[key][0]=[requested[i].split("-")]
            requested.pop(0)
            break

        for i in range(len(keys)):
            dictionary[key][1]=[not_wanted[i]]
            not_wanted.pop(0)
            break

    teams_split=[0]*len(keys)
    teams=[0]*len(keys)

    names=[]
    for i in range(len(keys)):
        names.append(dictionary_indices[i])

    for i in range(len(keys)):
        teams_split[i]=[0]
        teams[i]=[0]

    for i in range(len(teams_split)):
        for j in range(len(teams_split[i])):
            for k in range(len(names)):
                teams_split[i][j]=names[k]
                names.pop(k)
                break

    for i in range(len(teams_split)):
        teams[i]="-".join(teams_split[i])


    return [dictionary,teams_split,teams,dictionary_cost_features]            

"""
The cost function utilizes the above mentioned dictionaries in the following ways:
a. In case of incorrect team size, the respective member increments the cost by 1.
b. In case of incorrect team members, the respective member increments the cost by number of people that wasn't assigned correctly.
c. In case of unwanted team members, the respective member increments the cost by 2.
"""

def cost(t):

    for i in range(len(keys)):
        dictionary_cost[keys[i]]=0
        dictionary_cost_features[keys[i]]=[0,[]]

    for i in range(len(t)):
        for j in range(len(t[i])):
            dictionary_cost_features[t[i][j]][1].append(t[i])       #Gives teammates
            dictionary_cost_features[t[i][j]][0]=len(t[i])          #Gives length of team

            
    for key,value in dictionary_cost_features.items():
        if len(dictionary[key][0][0])!=dictionary_cost_features[key][0]:
            dictionary_cost[key]=1                                                      #If requested team size is not same

    
    for key,value in dictionary_cost_features.items():        
        for i in range(len(dictionary[key][0])):
            for k in range(len(dictionary[key][0][i])):
                if dictionary[key][0][i][k] not in dictionary_cost_features[key][1][0] and dictionary[key][0][i][k]!="xxx" and dictionary[key][0][i][k]!="zzz":
                    dictionary_cost[key]+=1                                                                 #If requested team members are not assigned
    
    for key,value in dictionary_cost_features.items():  
        for j in range(len(dictionary[key][1])):      
            if str(dictionary[key][1][j]) in dictionary_cost_features[key][1][0]:
                dictionary_cost[key]+=2                                                                         #If teammates not needed are assigned
          
    cost=0
    costs=[]
    for key,value in dictionary_cost.items():
        cost=cost+value
        costs.append(value)

    return [cost,costs,dictionary_cost_features,dictionary_cost]

"""
We generate the successors using the itertools library in Python.
We used the module to generate the combinations of several team members and we recursively call them our 1st state space is generated
"""

def successors(teams_split):
    team=copy.deepcopy(teams_split)				#Create a deepcopy of teams_split to form all permutations
    team_list=[]
    """Reference for next 2 lines of code is: https://www.geeksforgeeks.org/python-itertools-combinations-function/"""
    comb=combinations(teams_split,2)			#Make pairs of 2 people in a team
    y = ["+".join(i) for i in comb]				#Join with + to seperate groups from individuals
    """Reference for above 2 lines of code is: https://www.geeksforgeeks.org/python-itertools-combinations-function/"""
    x=[]


    for i in range(len(y)):
        x.append(y[i].split("-"))


    for i in range(len(y)):
        individual=y[i].split("+")
        for j in individual:
            if j in team:
                del team[team.index(j)]
        if y[i] in team:
            del team[team.index(y[i])]      
        

        team_list.append(team)
        team=copy.deepcopy(teams_split)
        team_list[i].append((y[i]))


    for i in range(len(team_list)):
        for j in team_list[i]:
            if '+' in j:
                team_list[i].append(str("-".join(j.split("+"))))		#Remove the plus and replace it with "-"
                team_list[i].pop(team_list[i].index(j))


    final_team=[]
    invalid=[]

    for i in range(len(team_list)):
        for j in team_list[i]:
            #print(j)
            x=j.split("-")
            if(len(x)<=3):								#Restrict team size to maximum of 3
                if team_list[i] not in final_team:
                    final_team.append(team_list[i])
            else:
                invalid.append(team_list[i])


    for i in range(len(invalid)):
        for j in final_team:
            if j==invalid[i]:
                final_team.pop(final_team.index(j))		#Remove all invalid teams of size>3.


    return final_team       

#As the input to the cost function is a 2-D array, we have to convert it into a 1-D array to perform our operations. Hence we have also defined a two_d_list function.

def two_d_list(array):
    for i in range(len(array)):
        array[i]=array[i].split("-")

    return array    

"""In the solve function we first generate the cost of the initial state. After that we try to generate states(teams) which will reduce the cost.
As there is not optimal solution for this problem, we have implemented a Local Search until we find the minimum value in a group of states. 
After sometime, if the system doesn't find an optimal state and if maximum team members reaches a limit of 3, it will get halted.
"""
def solver(input_file):

    check_cost=[]
    ip=initialize(input_file)			#Generate Initial State	
    complaints=cost(ip[1])				#Evaluate cost for intial state
    optimal_cost=complaints[0]			#Contains initial lowest cost
    teams=[]
    for i in range(len(ip[1])):			
        for j in range(len(ip[1][i])):
            teams.append(ip[1][i][j])


    output_dict={"assigned-groups":teams,"total-cost":complaints[0]}
    current_cost=complaints[0]
    
    yield({"assigned-groups":teams,"total-cost":complaints[0]})				#yield cost for initial state
    
    l=len(output_dict["assigned-groups"])

    while(current_cost>0 and len(output_dict["assigned-groups"])>(len(dictionary)//3)+1):	#We keep the while running till maximum of 3 members can be in a team

        succ_list=successors(teams)							#Generate Successor functions
        index_list=[]
        check_cost=[]
        dict_indices={}

        for i in range(len(succ_list)):
            c=cost(two_d_list(succ_list[i]))
            index_list.append(i)
            check_cost.append(c[0])							#Generate the cost for all successor states
        
        cost_dict={}
        
        zip_iterator = zip(index_list, succ_list)
        dict_indices = dict(zip_iterator)					#Map the cost to corresponding state through this dictionary.


        zip_iterator = zip(index_list, check_cost)
        cost_dict = dict(zip_iterator)						#Store the costs in a a dictionary to track the optimal one.

        optimal_cost=cost_dict[0]
        
        for key,value in cost_dict.items():
            if cost_dict[key]<optimal_cost:
                optimal_cost=cost_dict[key]					#Find the optimal cost
        

        for key,value in cost_dict.items():
            if optimal_cost==value:
                optimal_state=dict_indices[key]				#Find the optimal state


        teams=[0]*len(optimal_state)
        for i in range(len(optimal_state)):
            teams[i]="-".join(optimal_state[i])

        final_cost=cost(optimal_state)						#Calculate the correct cost
            
        if(optimal_cost>=current_cost):						#If current optimal state has greater than equal to previous optimal cost, then create successors of the current state.
            optimal_cost=current_cost
            succ_list=successors(teams)

        for key,value in output_dict.items():
            output_dict["total-cost"]=optimal_cost    		#Assign the optimal cost to "total cost" key.
                
        output_dict={"assigned-groups" : teams,"total-cost" : optimal_cost}

        if l==len(output_dict["assigned-groups"]):			#In case the length of number of teams of previous assigned group is the same as that of current one, recall the successor function.
            optimal_cost=current_cost
            succ_list=successors(teams)
        
        elif(l>len(output_dict["assigned-groups"])):
        	l=len(output_dict["assigned-groups"])
        	if optimal_cost==current_cost:					#In case the optimal cost is equal to the previous current cost, regenerate the successor function.
        		optimal_cost=current_cost
        		succ_list=successors(teams)
        	else:
        		current_cost=optimal_cost
        		yield({"assigned-groups" : teams,"total-cost" : optimal_cost})		#Yield the optimal solution.
     

    """
    1. This function should take the name of a .txt input file in the format indicated in the assignment.
    2. It should return a dictionary with the following keys:
        - "assigned-groups" : a list of groups assigned by the program, each consisting of usernames separated by hyphens
        - "total-cost" : total cost (number of complaints) in the group assignment
    3. Do not add any extra parameters to the solver() function, or it will break our grading and testing code.
    4. Please do not use any global variables, as it may cause the testing code to fail.
    5. To handle the fact that some problems may take longer than others, and you don't know ahead of time how
       much time it will take to find the best solution, you can compute a series of solutions and then
       call "yield" to return that preliminary solution. Your program can continue yielding multiple times;
       our test program will take the last answer you 'yielded' once time expired.
    """

    # Simple example. First we yield a quick solution
    """
    yield({"assigned-groups": ["vibvats-djcran-zkachwal", "shah12", "vrmath"],
               "total-cost" : 12})

    # Then we think a while and return another solution:
    time.sleep(10)
    yield({"assigned-groups": ["vibvats-djcran-zkachwal", "shah12-vrmath"],
               "total-cost" : 10})

    # This solution will never befound, but that's ok; program will be killed eventually by the
    #  test script.
    while True:
        pass
    
    yield({"assigned-groups": ["vibvats-djcran", "zkachwal-shah12-vrmath"],
               "total-cost" : 9})
    """

if __name__ == "__main__":
    if(len(sys.argv) != 2):
        raise(Exception("Error: expected an input filename"))

    for result in solver(sys.argv[1]):
        print("----- Latest solution:\n" + "\n".join(result["assigned-groups"]))
        print("\nAssignment cost: %d \n" % result["total-cost"])
    