import sys, json
from collections import Counter

vals={}
minScore=0.18

# threshold from lead_vs_nonleads.py
ratio_of_max_screen_threshold=0.748

demo={}
genders={}

eths=["black", "east_asian", "hispanic_latino", "south_asian", "white", "other"]

def get_gender_for_year(gender_json, year):
	ordered_years=sorted(list(gender_json.keys()))
	if len(ordered_years) == 1:
		return gender_json[ordered_years[0]].split("#")

	for idx, thisyear in enumerate(ordered_years):

		# if we've reach the last value, return that
		if idx == len(ordered_years)-1:
			return gender_json[thisyear].split("#")

		# otherwise check if the target year is between the current year and the next one in the list
		nextyear=ordered_years[idx+1]
		if year >= thisyear and year < nextyear:
			return gender_json[thisyear].split("#")


def str2int(x):
	return {int(k):v for k,v in x.items()}
	
def read_gender(filename):
	with open(filename) as file:
		for line in file:
			cols=line.rstrip().split("\t")
			idd=cols[0]
			gender_json=json.loads(cols[2], object_hook=str2int)
			genders[idd]=gender_json




def read_demo(filename):

	mapper={}
	mapper["Black (b)"]="black"
	mapper["East Asian (e)"]="east_asian"
	mapper["Hispanic/Latino (h)"]="hispanic_latino"
	mapper["None of the above (n)"]="other"
	mapper["South Asian/Indian (s)"]="south_asian"
	mapper["White (w)"]="white"


	with open(filename) as file:
		for line in file:
			cols=line.rstrip().split("\t")
			idd=cols[0]

			# if cols[13] == NA, then at least two perceptions are that this actor is not a real person and we should skip
			if cols[13] != "NA":

				demo[idd]={}
				for val in eths:
					demo[idd][val]=0
				for col in cols[1:11]:
					for val in col.split("#"):
						if val.rstrip() != "This is not a real person - e.g., cartoon character or animal (c)":
							demo[idd][mapper[val]]+=1





def read_meta(filename):
	with open(filename) as file:
		for line in file:
			cols=line.rstrip().split("\t")
			idd=cols[0]
			year=int(cols[3])
			vals[idd]=1

def read_fps(filename):

	with open(filename) as file:
		cols=file.readline().split("\t")

		frames=int(cols[1])
		fps=float(cols[4])
		return frames, fps

def read_recog(filename):
	recog={}
	with open(filename) as file:
		for line in file:
			cols=line.rstrip().split("\t")
			trackno=cols[0]

			actor, score=cols[3].split(" ")[0].split(":")
			recog[trackno]=actor, float(score)

	return recog


def get_tracks(filename):
	tracks={}
	tracks2fno={}
	with open(filename) as file:
		for line in file:
			cols=line.rstrip().split("\t")
			trackno=cols[0]
			fno=cols[1]
			faceno=cols[2]
			x1=max(0,int(cols[3]))
			y1=max(0,int(cols[4]))
			x2=int(cols[5])
			y2=int(cols[6])

			if fno not in tracks:
				tracks[fno]=[]
			tracks[fno].append((trackno, faceno, x1,y1,x2,y2))

			if trackno not in tracks2fno:
				tracks2fno[trackno]=[]
			tracks2fno[trackno].append(fno)
	return tracks, tracks2fno


def proc(dataFolder):

	for mov_id in vals:

		trackFile="%s/%s/%s.tracks.txt" % (dataFolder,mov_id,mov_id)
		recogFile="%s/%s/%s.recog.txt" % (dataFolder,mov_id,mov_id)
		fpsFile="%s/%s/%s.fps.txt" % (dataFolder,mov_id,mov_id)

		year=vals[mov_id]

		recog=read_recog(recogFile)
		tracks, tracks2fno=get_tracks(trackFile)
		total_frames, fps=read_fps(fpsFile)

		counts=Counter()

		for trackno in recog:
			actor, score=recog[trackno]
			if score >= 0.18:

				counts[actor]+=len(tracks2fno[trackno])


		max_actor_frames=counts.most_common()[0][1]

		lead_or_nonlead={}
		for actor in counts:
			ratio_of_max_screen=counts[actor]/max_actor_frames

			if ratio_of_max_screen > ratio_of_max_screen_threshold:
				lead_or_nonlead[actor]="lead"
			else:
				lead_or_nonlead[actor]="nonlead"


		for category in ["lead", "nonlead"]:

			actors={}
			eth_counts={}
			gender_counts={"male":{}, "female":{}}
			
			for eth in eths:
				eth_counts[eth]={}


			face_frames={}

			total_num_addressable_faces_in_frames=0
			total_num_addressable_faces_in_frames_gender=0

			for trackno in recog:
				for fno in tracks2fno[trackno]:
					face_frames[fno]=1


			eth_face_track_counts=Counter()
			gender_face_track_counts=Counter()


			for trackno in recog:

				actor, score=recog[trackno]

				if score > minScore:

					if lead_or_nonlead[actor] != category:
						continue

					if actor in demo:

						total_num_addressable_faces_in_frames+=len(tracks2fno[trackno])
		
						for dem in demo[actor]:

							percent=demo[actor][dem]/10

							eth_face_track_counts[dem]+=len(tracks2fno[trackno]) * percent

							for fno in tracks2fno[trackno]:							
								eth_counts[dem][fno]=percent

					if actor in genders:

						total_num_addressable_faces_in_frames_gender+=len(tracks2fno[trackno])
		
						for gender in get_gender_for_year(genders[actor], year):

							if gender not in gender_counts:
								gender_counts[gender]={}

							gender_face_track_counts[gender]+=len(tracks2fno[trackno])

							for fno in tracks2fno[trackno]:
								gender_counts[gender][fno]=1


			for dem in eth_counts:
				addressable_ratio=None
				if total_num_addressable_faces_in_frames > 0:
					addressable_ratio=eth_face_track_counts[dem]/total_num_addressable_faces_in_frames
				print("%s\t%s\t%s\t%s" % (mov_id, category, dem, addressable_ratio))
				sys.stdout.flush()

			for gender in gender_counts:
				addressable_ratio=None
				if total_num_addressable_faces_in_frames_gender > 0:
					addressable_ratio=gender_face_track_counts[gender]/total_num_addressable_faces_in_frames_gender
				print("%s\t%s\t%s\t%s" % (mov_id, category, gender, addressable_ratio))
				sys.stdout.flush()

	


metaFile=sys.argv[1]
demoFile=sys.argv[2]
genderFile=sys.argv[3]
dataFolder=sys.argv[4]
read_meta(metaFile)
read_demo(demoFile)
read_gender(genderFile)
proc(dataFolder)