from rfvision.models.detectors3d.category_ppf.category_ppf import CategoryPPF
from rfvision.models.detectors3d.category_ppf.utils.utils import pc_downsample, estimate_normals, backproject
from rfvision.models.detectors3d.category_ppf.utils.eval_utils import compute_degree_cm_mAP
from rflib.runner import load_checkpoint
import os
import glob
from tqdm import tqdm
import pickle
import cv2
import numpy as np
import torch

def get_nocs_pred(data_root):
    log_dir = os.path.join(data_root, 'real_test_20210511T2129')
    result_pkl_list = glob.glob(os.path.join(log_dir, 'results_*.pkl'))
    result_pkl_list = sorted(result_pkl_list)[:]

    final_results = []
    for pkl_path in tqdm(result_pkl_list):
        with open(pkl_path, 'rb') as f:
            result = pickle.load(f)
            if not 'gt_handle_visibility' in result:
                result['gt_handle_visibility'] = np.ones_like(result['gt_class_ids'])
                print('can\'t find gt_handle_visibility in the pkl.')
            else:
                assert len(result['gt_handle_visibility']) == len(result['gt_class_ids']), "{} {}".format(
                    result['gt_handle_visibility'], result['gt_class_ids'])

        if type(result) is list:
            final_results += result
        elif type(result) is dict:
            final_results.append(result)
        else:
            assert False
    return final_results

def get_model_for_all_categories(ppf_work_dir, epoch_id=200):
    model_for_all_categories = {}
    for i in range(1, 7):
        m = CategoryPPF(category=i).cuda().eval()
        checkpoint_path= os.path.join(ppf_work_dir, f'category{str(i)}', f'epoch_{str(epoch_id)}.pth')
        load_checkpoint(m, checkpoint_path)
        model_for_all_categories[i] = m
    return model_for_all_categories

synset_names = ['BG', #0
                'bottle', #1
                'bowl', #2
                'camera', #3
                'can',  #4
                'laptop',#5
                'mug'#6
                ]

synset_names_inv = dict([(k, v) for v, k in enumerate(synset_names)])

if __name__ == '__main__':
    data_root = '/disk1/data/ppf_dataset'
    final_results = get_nocs_pred(data_root)
    ppf_models = get_model_for_all_categories(ppf_work_dir='/home/hanyang/rfvision/work_dir/category_ppf',
                                              epoch_id=200)
    intrinsics = np.float32([[591.0125, 0, 320],[0, 590.16775, 240],[0, 0, 1]])
    out_dir = './'
    final_preds = {}
    ppf_model = ppf_models[1]
    for i , res in enumerate(tqdm(final_results)):
        res['image_path'] = data_root + res['image_path'][4:]
        img = cv2.imread(res['image_path'] + '_color.png')[:, :, ::-1]
        depth = cv2.imread(res['image_path'] + '_depth.png', -1)

        draw_img = img.copy()
        draw_image_bbox = img.copy()
        bboxs = res['pred_bboxes']
        masks = res['pred_masks']
        RTs = res['pred_RTs']
        scales = res['pred_scales']
        cls_ids = res['pred_class_ids']
        gt_RTs = res['gt_RTs']
        gt_cls_ids = res['gt_class_ids']

        ppf_pred_RT = []
        ppf_pred_scale = []

        for i, bbox in enumerate(bboxs):
            cv2.rectangle(draw_img,
                        (bbox[1], bbox[0]),
                        (bbox[3], bbox[2]),
                        (255, 0, 0), 2)
            draw_img[masks[:, :, i]] = np.array([0, 255, 0])
            cls_id = cls_ids[i]
            # ppf_model = ppf_models[cls_id]

            pc, idxs = backproject(depth, intrinsics, masks[:, :, i])
            pc /= 1000
            # augment
            # pc = pc + np.clip(cfg.res / 4 * np.random.randn(*pc.shape), -cfg.res / 2, cfg.res / 2)

            pc[:, 0] = -pc[:, 0]
            pc[:, 1] = -pc[:, 1]
            pc = pc_downsample(pc, 0.004)
            pc = pc.astype('float32')
            # high_res_indices = ME.utils.sparse_quantize(np.ascontiguousarray(pc), return_index=True, quantization_size=cfg.res)[1]
            # pc = pc[high_res_indices].astype(np.float32)
            pc_normal = estimate_normals(pc, 60).astype(np.float32)

            pcs = torch.from_numpy(pc[None]).cuda()
            pc_normals = torch.from_numpy(pc_normal[None]).cuda()
            ppf_preds = ppf_model.forward_test(pcs, pc_normals)
            ppf_pred_RT.append(ppf_preds['pred_RT'])
            ppf_pred_scale.append(ppf_preds['pred_scale'])

        ppf_pred_RT = np.array(ppf_pred_RT)
        ppf_pred_scale = np.array(ppf_pred_scale)

        final_results[i]['pred_RTs'] = ppf_pred_RT
        final_results[i]['pred_scales'] = ppf_pred_scale
    compute_degree_cm_mAP(final_results, synset_names, out_dir + '_plots',
                          degree_thresholds=[5, 10, 15],  # range(0, 61, 1),
                          shift_thresholds=[5, 10, 15],  # np.linspace(0, 1, 31)*15,
                          iou_3d_thresholds=np.linspace(0, 1, 101),
                          iou_pose_thres=0.1,
                          use_matches_for_pose=True)


# USAGE
# python detect_mask_video.py

# import the necessary packages
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.models import load_model
from imutils.video import VideoStream
from playsound import playsound
import numpy as np
import argparse
import imutils
import time
import cv2
import RPi.GPIO as GPIO #piny do diod led
import jetson.inference
import jetson.utils #biblioteka do obsługi kamery
import os
from plyer import notification

def wykryj_i_oszacuj_maske(klatka, faceNet, maskNet):
	# Wyciągnięcie rozmiarów klatki i storzenie z nich obiektu BLOB
	# Obiekty BLOB sluza do przechowywania duzych obiektow multimedialnych, np. obrazow
	(h, w) = klatka.shape[:2]
	blob = cv2.dnn.blobFromImage(klatka, 1.0, (300, 300),
		(104.0, 177.0, 123.0)) #

	# Uzycie sieci "faceNet" by wykryc twarze znajdujace sie na ujeciu z kamery
	faceNet.setInput(blob)
	wykryte_twarze = faceNet.forward()

	# inicjalizacja tablicy twarzy, ich lokacji i prawdopodobienst na maski
	twarze = []
	locs = []
	prawdopodobienstwo_maski = []

	# iterowanie po wykrytych twarzach
	for i in range(0, wykryte_twarze.shape[2]):
		# pobranie wartosci prawdopodoienstw. Prawdopodobienstwo mowi nam jaka jest szansa, ze
		# algorytm sie nie pomylil
		prawdopodobienstwo = wykryte_twarze[0, 0, i, 2]

		# Usuniece twarzy ktore maja prawdopodobienstwo mniejsze od min
		if prawdopodobienstwo > args["prawdopodobienstwo"]:
			# obliczenie koordynatow twarzy na ujeciu
			box = wykryte_twarze[0, 0, i, 3:7] * np.array([w, h, w, h])
			(startX, startY, endX, endY) = box.astype("int")

			# upewnienie sie, ze obliczone koordynaty znajduja sie na ujeciu.
			(startX, startY) = (max(0, startX), max(0, startY))
			(endX, endY) = (min(w - 1, endX), min(h - 1, endY))

			# Wyciagniecie ROI, czyli interesujacego nas obszaru i zamienienie go
			# z BGR na RGB
			# Dodatkowo zmieniamy rozmiary wykrytej twarzy na 224x224
			twarz = klatka[startY:endY, startX:endX]
			twarz = cv2.cvtColor(twarz, cv2.COLOR_BGR2RGB)
			twarz = cv2.resize(twarz, (224, 224))
			twarz = img_to_array(twarz)
			twarz = preprocess_input(twarz)

			# Dodanie twarzy i jej koordynatow do odpowiednich tablic
			twarze.append(twarz)
			locs.append((startX, startY, endX, endY))

	# Do kolejnego etapu przechodzimy tylko, jezeli wykrylismy min 1 twarz
	if len(twarze) > 0:
		# Na wykrytych twarzach przeprowadzamy kolejne oszacowanie, tym razem
		# sprawdzamy, z jakim prawdopodobienstwem na twarzy znajduje sie maska
		twarze = np.array(twarze, dtype="float32")
		prawdopodobienstwo_maski = maskNet.predict(twarze, batch_size=32)

	Adafruit_NeoPixel.setPixelColor(i, LED.Color(0, 0, 0)); #Wygaszenie diody gdy nie ma twarzy w kadrze
	Adafruit_NeoPixel.show();

	# zwracamy krotke (twarz, lokacja)
	return (locs, prawdopodobienstwo_maski)

#-----------------------------------------------
#Tutaj znajduje sie kod zwiazany z dioda, prezentujaca wykrycie maseczki
#LED variables
#Przygotowanie diod
def LED_PIN D1
def PWR_pin D2
def LED_COUNT 1
#Funkcja z biblioteki <Adafruit_NeoPixel.h>
Adafruit_NeoPixel = LED(LED_COUNT, LED_PIN, NEO_GRB + NEO_KHZ800);

# Przygotowujemy zmienne i podajemy lokalizacje modeli do wykrywania twarzy i masek
ap = argparse.ArgumentParser()
ap.add_argument("-f", "--twarz", type=str,
	default="wykrywacz_twarzy",
	help="sciezka do folderu z modelem do wykrywania twarzy")
ap.add_argument("-m", "--model", type=str,
	default="mask_detector.model",
	help="sciezka do folderu z modelem do wykrywania twarzy")
ap.add_argument("-c", "--prawdopodobienstwo", type=float, default=0.5,
	help="minimalne prawdopodobienstwo do filtrowania twarzy")
args = vars(ap.parse_args())

# Wczytanie modelu do wykrywania twarzy z dysku
print("[INFO] ladowanie modelu do wykrywania twarzy...")
prototxtPath = os.path.sep.join([args["twarz"], "deploy.prototxt"])
weightsPath = os.path.sep.join([args["twarz"],
	"res10_300x300_ssd_iter_140000.caffemodel"])
faceNet = cv2.dnn.readNet(prototxtPath, weightsPath)

# Wczytanie modelu do wykrywania maseczek z dysku
print("[INFO] ladowanie modelu do wykrywania maseczek...")
maskNet = load_model(args["model"])

#-----------------------------------------------------
# inicjalizowanie strumienia z kamery
print("[INFO] startowanie przesylania obrazu...")
camera = jetson.utils.gstCamera(1280, 720, "0")
vs = camera.CaptureRGBA(zeroCopy=1)
time.sleep(2.0)
#-----------------------------------------------------

# iterowanie po klatkach z strumienia z kamery
while True:
	# Pobranie klatki obrazu i zmiana jego rozmiaru
	klatka = vs.read()
	klatka = imutils.resize(klatka, width=400)

	# Wykryj twarze na ujeciu i oszacuj czy sa w maseczc
	(locs, prawdopodobienstwo_maski) = wykryj_i_oszacuj_maske(klatka, faceNet, maskNet)

	# Petla iterujaca po wszystkich wykrytych lokacjach twarzy
	for (box, pred) in zip(locs, prawdopodobienstwo_maski):
		# unpack the bounding box and predictions
		(startX, startY, endX, endY) = box
		(maska, bezMaski) = pred

		# Ustalenie etykiety i koloru obramowania wykrytych obiektow na obrazie
		label = "Maska" if maska > bezMaski else "Bez Maski"
		color = (0, 255, 0) if label == "Maska" else (0, 0, 255)

		# Dodanie prawdopodobienstwa do etykiety obiektu
		label = "{}: {:.2f}%".format(label, max(maska, bezMaski) * 100)

		# Wystwietlenie etykiety i obramowania na kadrze
		cv2.putText(klatka, label, (startX, startY - 10),
			cv2.FONT_HERSHEY_SIMPLEX, 0.45, color, 2)
		cv2.rectangle(klatka, (startX, startY), (endX, endY), color, 2)

		# Tutaj reagujemy na maske lub jej brak
		if maska < bezMaski:
			Adafruit_NeoPixel.setPixelColor(i, pixels.Color(0, 150, 0)); #Zapalenie diody na czerwono
		else
		    Adafruit_NeoPixel.setPixelColor(i, pixels.Color(150, 0, 0)); #Zapalenie diody na zielono

		Adafruit_NeoPixel.show();

	# wyswietlenie klatki
	cv2.imshow("klatka", klatka)
	key = cv2.waitKey(1) & 0xFF

	# Wylaczenie programu przyciskiem q
	if key == ord("q"):
		break

# "Sprzatanie", wylaczenie strumienia z kamery i zniszczenie wszystkich okien
cv2.destroyAllWindows()
vs.stop()


# Matt Bonnyman 24 July 2018

# astroconda packages
import argparse
import astropy.units as u
from astroplan import download_IERS_A
from astropy.table import Table, Column
from astropy.time import Time
import importlib
import numpy as np
import textwrap

# gemini packages
from sb import sb
import convert_conditions
from condition_table import condition_table
from dates import getdates
import make_plot
import printer
import program_table
from time_table import time_table
import timing_windows
import weights
from wind_table import wind_table
from observation_table import observation_table
from sun_table import sun_table
from moon_table import moon_table
from instrument_table import instrument_table
import select_obs as select_obs
from observer_site import getsite
import convert_conditions as convertcond

def getobsindex(obs_id, obs):
    """
    Retrieve the index for the row of an observation the observation data table 'obs'.

    Parameters
    ----------
    obs_id : string
        Uniqie observation identifier

    obs : '~astropy.table.Table'
        Observation data table or subset of table.

    Returns
    -------
    integer of row index corresponding to selected observation.
    """

    ii = np.where(obs['obs_id'] == obs_id)[0][:]
    if len(ii) == 1:
        return ii[0]
    elif len(ii) == 0:
        raise ValueError(' \'{}\' not found.'.format(obs_id))
    elif len(ii) > 1:
        raise ValueError('Found ' + str(len(ii)) + ' observations with identifier ' + obs_id)
    else:
        raise ValueError('Something went wrong.')

def weightplotmode(site, timetable, sun, moon, obs, progs, targets, skycond, wind, wra):
    """
    Select and plot the weighting function of any observation on the current date.  When this
    function is run separately from the Gemini Queue Planning Tool, it will not generate a plan,
    and will therefore not be able to print the plan table, and will not require the 'plan' or 'targets'
    variables.

    Parameters
    ----------
    site : '~astroplan.Observer'
        Observatory site information.

    timetable : '~astropy.table.Table'
        time information table created by time_table.py

    sun : '~astropy.table.Table'
        sun information table created by sun_table.py

    moon : '~astropy.table.Table'
        moon information table created by moon_table.py

    obs : '~astropy.table.Table'
        observation information table created by observation_table.py

    progs : '~astropy.table.Table'
        program status information table created by observation_table.py

    targets : '~astropy.table.Table'
        target position data table for observing period created by target_table.py

    skycond : '~astropy.table.Table'
        sky conditions information table created by condition_table.py

    wind : '~astropy.table.Table'
        wind conditions information table created by wind_table.py

    wra : numpy float array
        Right ascension distribution weights computed from full set of observations.s

    """
    verbose = False

    aprint = '\t{0:<20s}{1}'  # print two strings
    fprint = '\t{0:5s}{1:40}{2:40s}'  # format options menu

    print('\n\n\t---------- Weight function plotting mode ----------\n')
    print(aprint.format('Plan date:', timetable['date'][0]))
    while True:
        # -- Weight function plotting mode menu --
        print('\n\tOptions:')
        print('\t--------')
        print(fprint.format('1.', 'See list of available observations', '-'))
        print(fprint.format('2.', 'Conditions (iq,cc,wv)',
                            '({}, {}, {})'.format(str(skycond['iq'][0]), str(skycond['cc'][0]), str(skycond['wv'][0]))))
        print(fprint.format('3.', 'Wind conditions (dir, vel)',
                            '({}, {})'.format(wind['dir'].quantity[0], wind['vel'].quantity[0])))
        print(fprint.format('x,q', 'Exit', ''))

        rawinput = input('\n Select option or provide an observation identifier: ')
        userinput = rawinput.strip()

        if userinput == '1':
            [print(line) for line in printer.listobs(obs=Table(obs['obs_id', 'prog_ref', 'group', 'target']))]
            continue

        elif userinput == '2':
            condinput = input(' Input new iq, cc, wv percentiles(eg. \'20 50 Any\'): ')
            tempconds = condinput.split(' ')
            if len(tempconds) != 3:
                print(' Did not receive 3 values. No changes were made.')
                continue
            else:
                try:
                    iq, cc, wv = convertcond.inputcond(iq=tempconds[0].strip("' "),
                                                       cc=tempconds[1].strip(), wv=tempconds[2].strip("' "))
                    skycond = condition_table(size=len(timetable['utc'].data[0]), iq=iq, cc=cc, wv=wv)
                except ValueError:
                    print(' ValueError: Could not set new conditions. Changes not made.')
                    continue

        elif userinput == '3':
            windinput = input(' Input new direction(deg), velocity(m/s) (eg. \'330 5\'): ')
            tempwind = windinput.split(' ')
            if len(tempwind) != 2:
                print(' Did not receive 2 values. No changes were made.')
                continue
            else:
                try:
                    wind = wind_table(size=len(timetable['utc'][0]), direction=float(tempwind[0].strip("' ")),
                                      velocity=float(tempwind[1].strip("' ")), site_name=site.name)
                except ValueError:
                    print(' ValueError: Could not set new wind conditions.')
                    continue

        elif userinput.lower() == 'quit' or userinput.lower() == 'exit' \
                or userinput.lower() == 'q' or userinput.lower() == 'x':
            break

        else:
            try:  # retrieve table row number
                i = getobsindex(obs_id=userinput, obs=Table(obs))
            except ValueError as e:
                print(e)
                continue

            importlib.reload(make_plot)
            importlib.reload(weights)

            target = Table(targets[i])

            # ====== Compute visible sky brightnesses at targets ======
            target['vsb'] = Column([sb(mpa=moon['phase'].quantity[0],
                                       mdist=targets['mdist'].quantity[i],
                                       mZD=moon['ZD'].quantity[0],
                                       ZD=targets['ZD'].quantity[i],
                                       sZD=sun['ZD'].quantity[0],
                                       cc=skycond['cc'].data)])

            # ====== Convert vsb to sky background percentiles ======
            target['bg'] = Column([convert_conditions.sb_to_cond(sb=target['vsb'][0])])

            target['weight'] = Column(
                [weights.obsweight(
                    obs_id=obs['obs_id'][i],
                    ra=obs['ra'].quantity[i],
                    dec=obs['dec'].quantity[i],
                    iq=obs['iq'].data[i],
                    cc=obs['cc'].data[i],
                    bg=obs['bg'].data[i],
                    wv=obs['wv'].data[i],
                    elev_const=obs['elev_const'][i],
                    i_wins=target['i_wins'][0],
                    band=obs['band'].data[i],
                    user_prior=obs['user_prior'][i],
                    AM=target['AM'].data[0],
                    HA=target['HA'].quantity[0],
                    AZ=target['AZ'].quantity[0],
                    latitude=site.location.lat,
                    prog_comp=progs['prog_comp'].data[obs['i_prog'].data[i]],
                    obs_comp=obs['obs_comp'].quantity[i],
                    skyiq=skycond['iq'].data,
                    skycc=skycond['cc'].data,
                    skywv=skycond['wv'].data,
                    skybg=target['bg'].data[0],
                    winddir=wind['dir'].quantity,
                    windvel=wind['vel'].quantity,
                    wra=wra[i])])

            make_plot.weightcomponents(
                    obs_id=obs['obs_id'][i],
                    ra=obs['ra'].quantity[i],
                    dec=obs['dec'].quantity[i],
                    iq=obs['iq'].data[i],
                    cc=obs['cc'].data[i],
                    bg=obs['bg'].data[i],
                    wv=obs['wv'].data[i],
                    elev_const=obs['elev_const'][i],
                    i_wins=target['i_wins'][0],
                    band=obs['band'].data[i],
                    user_prior=obs['user_prior'][i],
                    AM=target['AM'].data[0],
                    HA=target['HA'].data[0] * u.hourangle,
                    AZ=target['AZ'].quantity[0],
                    latitude=site.location.lat,
                    prog_comp=progs['prog_comp'].data[obs['i_prog'].data[i]],
                    obs_comp=obs['obs_comp'].quantity[i],
                    skyiq=skycond['iq'].data,
                    skycc=skycond['cc'].data,
                    skywv=skycond['wv'].data,
                    skybg=target['bg'].data[0],
                    winddir=wind['dir'].quantity,
                    windvel=wind['vel'].quantity,
                    wra=wra[i],
                    localtimes = timetable['local'].data[0])

            make_plot.weightfunction(obs_id=userinput,
                                     local_time=timetable['local'].data[0],
                                     date=timetable['date'].data[0],
                                     weight=target['weight'].quantity[0])

            make_plot.skyconditions(skycond=skycond,
                                    local_time=timetable['local'].data[0],
                                    date=timetable['date'].data[0],
                                    bg=target['bg'].quantity[0], verbose=verbose)

            make_plot.vsb(vsb=target['vsb'][0], local_time=timetable['local'].data[0],date=timetable['date'].data[0],
                          obs_id=obs['obs_id'][i])

            make_plot.windconditions(wind=wind,
                                     local_time=timetable['local'].data[0],
                                     date=timetable['date'].data[0])
    return

def wfpt():

    parser = argparse.ArgumentParser(prog='wfpt.py',
                                     formatter_class=argparse.RawDescriptionHelpFormatter,
                                     description=textwrap.dedent('''
                                        Weight function plotting tool
        *****************************************************************************************************

            otfile                  OT catalog file name.

            prfile                  Gemini exechours program status file name.

            instcal                 Instrument calendar filename.

            -o   --observatory      Observatory site [DEFAULT='gemini_south']. Accepts the following:
                                    1. 'gemini_north' (or 'MK' for Mauna Kea)
                                    2. 'gemini_south' (or 'CP' for Cerro Pachon)

            -d   --date             Date 'YYYY-MM-DD' [DEFAULT=current].

            -dst --daylightsavings  Toggle daylight savings time [DEFAULT=False].

            -dt  --gridsize         Size of time-grid spacing [DEFAULT=0.1hr].

                                    Sky conditions:
            -i   --iq               Image quality constraint [DEFAULT=70].
            -c   --cc               Cloud cover constraint   [DEFAULT=50].
            -w   --wv               Water vapor constraint   [DEFAULT=Any].

                                    Wind conditions:
            -dir --direction        Wind direction [DEFAULT=270deg].
            -vel --velocity         Wind velocity [DEFAULT=10deg].

            -rw  --randwind         Random wind conditions (use mean and standard deviation of site):
                                        Cerro Pachon : dir=(330 +/- 20)deg, vel=(5 +/- 3)m/s
                                        Mauna Kea    : dir=(330 +/- 20)deg, vel=(5 +/- 3)m/s

            -u   --update           Download up-to-date IERS(International Earth Rotation and Reference Systems).

        *****************************************************************************************************
                                        '''))

    parser.add_argument(action='store',
                        dest='otfile')

    parser.add_argument(action='store',
                        dest='prfile')

    parser.add_argument(action='store',
                        dest='instfile')

    parser.add_argument('-o', '--observatory',
                        action='store',
                        default='gemini_south')

    parser.add_argument('-d', '--date',
                        action='store',
                        default=None)

    parser.add_argument('-dst', '--daylightsavings',
                        action='store_true',
                        dest='dst',
                        default=False)

    parser.add_argument('-dt', '--gridsize',
                        action='store',
                        default=0.1)

    parser.add_argument('-iq', '--iq',
                        default='70')

    parser.add_argument('-cc', '--cc',
                        default='50')

    parser.add_argument('-wv', '--wv',
                        default='Any')

    parser.add_argument('-dir', '--direction',
                        default=330)

    parser.add_argument('-vel', '--velocity',
                        default=5)

    parser.add_argument('-rw', '--randwind',
                        action='store_true',
                        default=False)

    parser.add_argument('-u', '--update',
                        action='store_true',
                        default=False)

    parse = parser.parse_args()
    otfile = parse.otfile
    prfile = parse.prfile
    instfile = parse.instfile
    iq, cc, wv = convertcond.inputcond(parse.iq, parse.cc, parse.wv)
    dst = parse.dst
    dir = parse.direction
    vel = parse.velocity
    randwind = parse.randwind

    if parse.update:  # download most recent International Earth Rotation and Reference Systems data
        download_IERS_A()

    verbose = False
    verbose_progress = True

    # Time grid spacing size in hours
    dt = float(parse.gridsize) * u.h
    if dt <= 0 * u.h:
        raise ValueError('Time grid spacing must be greater than 0.')

    # ====== Get Site Info ======
    if verbose_progress:
        print('...observatory site, time_zone, utc_to_local')
    # Create 'astroplan.Observer' object for observing site.
    # Get time-zone name (for use by pytz) and utc_to_local time difference.
    site, timezone_name, utc_to_local = getsite(site_name=parse.observatory,
                                                daylightsavings=dst)

    # ====== Plan start/end dates ======
    if verbose_progress:
        print('...scheduling period dates')
    # Check format of command line input dates.
    # Create 'astropy.time.core.Time' objects for start and end of plan period
    start, end = getdates(startdate=parse.date,
                          enddate=None,
                          utc_to_local=utc_to_local)

    # ====== Time data table for scheduling period ======
    if verbose_progress:
        print('...time data and grids')
    # Create 'astropy.table.Table' (one row for each day in plan period)
    # Stores time grids for UTC, local, lst.
    # Stores solar midnight, evening/morning nautical twilights
    timetable = time_table(site=site,
                           utc_to_local=utc_to_local,
                           dt=dt,
                           start=start,
                           end=end)

    # set start and end times as boundaries of scheduling period
    start = Time(timetable[0]['utc'][0])
    end = Time(timetable[-1]['utc'][-1])

    # ====== Sun data table for scheduling period ======
    if verbose_progress:
        print('...Sun data')
    # Create 'astropy.table.Table' (one row for each day in plan period)
    # Stores ra, dec at midnight on each night.
    # Stores azimuth angle, zenith distance angle, and hour angle
    # throughout scheduling period.
    sun = sun_table(latitude=site.location.lat,
                    solar_midnight=timetable['solar_midnight'].data,
                    lst=timetable['lst'].data)

    # ====== Moon data table for scheduling period ======
    if verbose_progress:
        print('...Moon data')
    # Create 'astropy.table.Table' (one row for each day in plan period)
    # Stores fraction illuminated, phase angle, ra, and dec at solar midnight
    # on each night.
    # Stores ra, dec, azimuth angle, zenith distance angle, hour angle, airmass
    # throughout scheduling period.
    moon = moon_table(site=site,
                      solar_midnight=timetable['solar_midnight'].data,
                      utc=timetable['utc'].data,
                      lst=timetable['lst'].data,
                      verbose=False)

    # ====== Assemble Observation Table ======
    if verbose_progress:
        print('...observations')
    # Create 'astropy.table.Table' (one observation per row)
    obs = observation_table(filename=otfile)

    # ====== Assemble Program Table ======
    if verbose_progress:
        print('...programs')
    # read columns from exechours_YYYYL.txt file into 'astropy.table.Table' object
    exechourtable = program_table.read_exechours(filename=prfile)

    # retrieve additional program information from the observation table (if available)
    proginfo = program_table.get_proginfo(exechourtable=exechourtable,
                                          prog_ref_obs=obs['prog_ref'].data,
                                          obs_id=obs['obs_id'].data,
                                          pi=obs['pi'].data,
                                          partner=obs['partner'].data,
                                          band=obs['band'].data,
                                          too_status=obs['too_status'].data)

    # For now, program activation(prog_start) and
    # deactivation(prog_end) times are are set to the scheduling period
    # boundaries.
    # All programs are set to active.
    # These inputs can be changed once the information is available.
    progs = program_table.programtable(gemprgid=proginfo['prog_ref'].data,
                                       partner=proginfo['partner'].data,
                                       pi=proginfo['pi'].data,
                                       prog_time=proginfo['prog_time'].data,
                                       alloc_time=proginfo['alloc_time'].data,
                                       partner_time=proginfo['partner_time'].data,
                                       active=np.full(len(proginfo), True),
                                       prog_start=np.full(len(proginfo), start),
                                       prog_end=np.full(len(proginfo), end),
                                       too_status=proginfo['too_status'].data,
                                       scirank=proginfo['scirank'].data,
                                       observations=proginfo['obs'].data)

    # Add an additional column in the observation table to hold the
    # program table (progs) row index of each observation's
    # corresponding program.
    obs['i_prog'] = select_obs.i_progs(gemprgid=progs['gemprgid'].data,
                                       prog_ref=obs['prog_ref'].data)

    # ====== Instrument configuration calendar table ======
    # Create 'astropy.table.Table' (one row per night in plan period)
    # Store date, available instruments, GMOS-FPU, GMOS-Disperser, F2-FPU
    if verbose_progress:
        print('...instrument calendar')
    instcal = instrument_table(filename=instfile,
                               dates=timetable['date'].data)

    # ====== Timing windows and target calendar ======
    if verbose_progress:
        print('...target calendar')
    targetcal = timing_windows.get_timing_windows(site=site,
                                                  timetable=timetable,
                                                  moon=moon,
                                                  obs=obs,
                                                  progs=progs,
                                                  instcal=instcal)

    # -- Conditions tables --
    skycond = condition_table(size=len(timetable['utc'][0]), iq=iq, cc=cc, wv=wv)
    wind = wind_table(size=len(timetable['utc'][0]), direction=dir, velocity=vel, site_name=site.name)
    if verbose:
        print('skycond\n',skycond)
        print('wind\n',wind)

    # ====== Get remaining available observations in nightly queue ======
    # target_cal[i_day] observations with remaining program time
    i_queue_cal = np.where(progs['prog_comp'].data[obs['i_prog'].data[targetcal[0]['i'].data]] < 1)[0][:]

    # target_cal[i_day] observations with remaining observation time
    i_queue_cal = i_queue_cal[np.where(obs['obs_comp'].data[targetcal[0]['i'].data[i_queue_cal]] < 1)[0][:]]

    # obs table rows of observations with remaining observation and program time
    i_queue_obs = targetcal[0]['i'].data[i_queue_cal]

    # all observations with remaining observation and program time in full queue.
    # Required for computing the distribution of remaining observation time.
    i_obs = np.where(progs['prog_comp'].data[obs['i_prog'].data] < 1)[0][:]
    i_obs = i_obs[np.where(obs['obs_comp'].data[i_obs] < 1)[0][:]]

    # ====== Compute observation time distribution (wra = right ascension weight) ======
    wra_all = weights.radist(ra=obs['ra'].quantity[i_obs],
                             tot_time=obs['tot_time'].quantity[i_obs],
                             obs_time=obs['obs_time'].quantity[i_obs])  # wra of all obs
    wra = wra_all[np.where([(i in i_obs) for i in i_queue_obs])[0][:]]  # wra of obs in tonight's queue

    # ====== tonight's queue from targetcal table ======
    # Create 'astropy.table.Table' for the active observations in tonight's queue
    targets = Table(targetcal[0][i_queue_cal])

    weightplotmode(site=site, timetable=timetable, sun=sun, moon=moon, obs=Table(obs[i_queue_obs]), progs=progs,
                   targets=targets, skycond=skycond, wind=wind, wra=wra)

    return

if __name__ == '__main__':
    wfpt()

import numpy,import astropy,from astropy


#!/usr/bin/python3

import argparse
import os
import random
import logging
from scipy import misc

import tensorflow as tf
import numpy as np
from PIL import Image

from unet import UNet, Discriminator
from scripts.image_manips import resize
import cv2

model_name = "matting"
TRAIN_IMAGES = 50

logging.basicConfig(level=logging.INFO)

# Parse Arguments
def parse_args():
    parser = argparse.ArgumentParser(description="Trains the unet")
    parser.add_argument("data", type=str,
        help="Path to a folder containing data to train")
    parser.add_argument("--lr", type=float, default=1.0,
        help="Learning rate used to optimize")
    parser.add_argument("--d_coeff", type=float, default=1.0,
        help="Discriminator loss coefficient")
    parser.add_argument("--gen_epoch", type=int, default=40,
        help="Number of training epochs")
    parser.add_argument("--disc_epoch", type=int, default=10,
        help="Number of training epochs")
    parser.add_argument("--adv_epoch", type=int, default=50,
        help="Number of training epochs")
    parser.add_argument("--batch_size", dest="batch_size", type=int, default=4,
        help="Size of the batches used in training")
    parser.add_argument('--checkpoint', type=int, default=None,
        help='Saved session checkpoint, -1 for latest.')
    parser.add_argument('--logdir', default="log/" + model_name,
        help='Directory where logs should be written.')
    return  parser.parse_args()

def apply_trimap(images, output, alpha):
    masked_output = []
    for channel in range(4):
        masked_output.append(output[:,:,:,channel])
        masked_output[channel] = tf.where(alpha < 0.25, images[:,:,:,channel], masked_output[channel])
        masked_output[channel] = tf.where(alpha > 0.75, images[:,:,:,channel], masked_output[channel])
        masked_output[channel] = masked_output[channel]
    masked_output = tf.stack(masked_output, 3)
    return masked_output

def main(args):
    input_path = os.path.join(args.data, "original")
    trimap_path = os.path.join(args.data, "original_depth")
    target_path = os.path.join(args.data, "bokeh")
    output_path = os.path.join(args.data, "output")
    print('input_path '+input_path)
    train_data_update_freq = args.batch_size
    test_data_update_freq = 50*args.batch_size
    sess_save_freq = 100*args.batch_size

    if not os.path.isdir(output_path):
        os.makedirs(output_path)

    if not os.path.isdir(args.logdir):
        os.makedirs(args.logdir)

    ids = [[int(i) for i in os.path.splitext(filename)[0].split('_')] for filename in os.listdir(input_path)]
    print("numimage "+str(len(ids)))
    np.random.shuffle(ids)
    split_point = int(round(0.85*len(ids))) #using 70% as training and 30% as Validation
    train_ids = tf.get_variable('train_ids', initializer=ids[0:split_point], trainable=False)
    valid_ids = tf.get_variable('valid_ids', initializer=ids[split_point:len(ids)], trainable=False)

    global_step = tf.get_variable('global_step', initializer=0, trainable=False)

    g_iter = int(args.gen_epoch * int(train_ids.shape[0]))
    d_iter = int(args.disc_epoch * int(train_ids.shape[0]))
    a_iter = int(args.adv_epoch * int(train_ids.shape[0]))
    n_iter = g_iter+d_iter+a_iter
    print('n_iter ' + str(n_iter))

    input_images = tf.placeholder(tf.float32, shape=[TRAIN_IMAGES, 480, 360, 4])
    target_images = tf.placeholder(tf.float32, shape=[TRAIN_IMAGES, 480, 360, 3])

    #alpha = target_images[:,:,:,3][..., np.newaxis]

    with tf.variable_scope("Gen"):
        gen = UNet(4,3)
        output = tf.sigmoid(gen(input_images))
        g_loss = tf.losses.mean_squared_error(target_images, output)
    with tf.variable_scope("Disc"):
        disc = Discriminator(4)
        d_real = disc(target_images)
        d_fake = disc(output)
        d_loss = tf.reduce_mean(tf.log(d_real) + tf.log(1-d_fake))

    a_loss = g_loss + args.d_coeff * d_loss

    g_loss_summary = tf.summary.scalar("g_loss", g_loss)
    d_loss_summary = tf.summary.scalar("d_loss", d_loss)
    a_loss_summary = tf.summary.scalar("a_loss", a_loss)

    summary_op = tf.summary.merge(
        [g_loss_summary, d_loss_summary, a_loss_summary])

    summary_image = tf.summary.image("result", output)

    g_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='Gen')
    d_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='Disc')

    g_optimizer = tf.train.AdadeltaOptimizer(args.lr).minimize(g_loss, global_step=global_step, var_list=g_vars)
    a_optimizer = tf.train.AdadeltaOptimizer(args.lr).minimize(a_loss, global_step=global_step, var_list=g_vars)
    d_optimizer = tf.train.AdadeltaOptimizer(args.lr).minimize(-d_loss, global_step=global_step, var_list=d_vars)

    init = tf.global_variables_initializer()
    sess = tf.Session()
    sess.run(init)

    train_writer = tf.summary.FileWriter(args.logdir + '/train')
    test_writer = tf.summary.FileWriter(args.logdir + '/test')
    saver = tf.train.Saver()
    if args.checkpoint is not None and os.path.exists(os.path.join(args.logdir, 'checkpoint')):
        if args.checkpoint == -1:#latest checkpoint
            saver.restore(sess, tf.train.latest_checkpoint(args.logdir))
        else:#Specified checkpoint
            saver.restore(sess, os.path.join(args.logdir, model_name+".ckpt-"+str(args.checkpoint)))
        print('Model restored to step ' + str(global_step.eval(sess)))

    train_ids = list(train_ids.eval(sess))
    valid_ids = list(valid_ids.eval(sess))

    def load_batch(batch_ids):
        images, targets = [], []
        test_data = np.zeros((TRAIN_IMAGES, 480, 360, 4))
        test_answ = np.zeros((TRAIN_IMAGES, int(480 * 1), int(360 * 1), 3))
        i = 0
        for names in os.listdir(input_path):
            if i < TRAIN_IMAGES:
                random_file = random.choice(os.listdir(input_path))
                name = os.path.basename(random_file.split(".")[0])
                input_filename = os.path.join(input_path, name + '.jpg')
                trimap_filename = os.path.join(trimap_path, name + '.png')
                target_filename = os.path.join(target_path, name + '.png')
                I = misc.imread(input_filename)
                I_depth = misc.imread(trimap_filename)

                if i == 20:
                    print('random_file: '+random_file)

                I = cv2.resize(I, (360, 480));
                I_depth = cv2.resize(I_depth, (360, 480));
                # Stacking the image together with its depth map
                I_temp = np.zeros((I.shape[0], I.shape[1], 4))
                I_temp[:, :, 0:3] = I
                I_temp[:, :, 3] = I_depth
                I = I_temp

                # Extracting random patch of width PATCH_WIDTH

                I = np.float32(I) / 255.0

                test_data[i, :] = I

                I = misc.imread(target_filename)
                I = cv2.resize(I, (int(360), int(480)));

                I = np.float32(I) / 255.0

                test_answ[i, :] = I
                i = i + 1

        return test_data, test_answ

    def test_step(batch_idx, summary_fct):
        batch_range = random.sample(train_ids, args.batch_size)

        images, targets = load_batch(batch_range)

        loss, demo, summary = sess.run([g_loss, summary_image, summary_fct], feed_dict={
            input_images: images,
            target_images: targets,
            })

        test_writer.add_summary(summary, batch_idx)
        test_writer.add_summary(demo, batch_idx)
        print("---------------------------------------------")
        print(output)
        image1_tensor= output[0]
        print(image1_tensor.shape)
        image1_numpy = tf.cast(image1_tensor, tf.uint8)
        print(image1_numpy)
        print('Validation Loss: {:.8f}'.format(loss))

    batch_idx = 0
    while batch_idx < n_iter:
        batch_idx = global_step.eval(sess) * args.batch_size
        print('batch_idx' + str(batch_idx))

        loss_fct = None
        label = None
        optimizers = []
        if batch_idx < g_iter:
            print('Gen train')
            loss_fct = g_loss
            summary_fct = g_loss_summary
            label = 'Gen train'
            optimizers = [g_optimizer]
        elif batch_idx < g_iter+d_iter:
            loss_fct = d_loss
            summary_fct = d_loss_summary
            label = 'Disc train'
            optimizers = [d_optimizer]
            print('Disc train')
        else:
            loss_fct = a_loss
            summary_fct = summary_op
            label = 'Adv train'
            optimizers = [a_optimizer]
            print('Adv train')

        batch_range = random.sample(train_ids, args.batch_size)
        images, targets = load_batch(batch_range)
        print('load batch done')
        loss, summary = sess.run([loss_fct, summary_fct] +  optimizers, feed_dict={
            input_images: images,
            target_images: targets})[0:2]
        print('load batch done 1')
        if batch_idx % train_data_update_freq == 0:
            print('{}: [{}/{} ({:.0f}%)]\tGen Loss: {:.8f}'.format(label, batch_idx, n_iter,
                100. * (batch_idx+1) / n_iter, loss))
            print('load batch done 2')
            train_writer.add_summary(summary, batch_idx)

        if batch_idx % test_data_update_freq == 0:
            print('load batch done 3')
            test_step(batch_idx, summary_fct)

        if batch_idx % sess_save_freq == 0:
            print('Saving model')
            saver.save(sess, os.path.join(args.logdir, model_name+".ckpt"), global_step=batch_idx)

    saver.save(sess, os.path.join(args.logdir, model_name+".ckpt"), global_step=batch_idx)

if __name__ == '__main__':
    args = parse_args()
    main(args)


# -*- coding: utf-8 -*-
# Copyright (c) 2017 - for information on the respective copyright owner
# see the NOTICE file and/or the repository https://github.com/boschresearch/statestream
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" Example of wrapping a C library function that accepts a C double array as
    input using the numpy.ctypeslib.
"""

import numpy as np
import numpy.ctypeslib as npct
from ctypes import c_int, c_float32

# Define some default variables for c-functions.
array_1d_float = npct.ndpointer(dtype=np.float32, ndim=1, flags="CONTIGUOUS")
array_1d_int = npct.ndpointer(dtype=np.int, ndim=1, flags="CONTIGUOUS")

# Load the compiled c-library, using numpy mechanisms.
libcd = npct.load_library("ccpp/lib/libchebbian", ".")

# Setup the return types and argument types.

void chebb_raw(float* src,
                float* tgt,
                float* upd,
                float tgt_eps,
                int agents,
                int src_X,
                int src_Y,
                int src_C,
                int tgt_X,
                int tgt_Y,
                int tgt_C,
                int rf_X,
                int rf_Y,
                int dil_X,
                int dil_Y,
                int ignore_border)

libcd.chebb_raw.restype = None
libcd.chebb_raw.argtypes = [array_1d_float,
                            array_1d_float,
                            array_1d_float,
                            c_float32,
                            c_int,
                            c_int,
                            c_int,
                            c_int,
                            c_int,
                            c_int,
                            c_int,
                            c_int,
                            c_int,
                            c_int,
                            c_int,
                            c_int]

# Function for fast blitting and color conversion of images.
def chebb_raw2D(source, target, update, eps=1e-2, dilation=(1,1), ignore_border=True):
    """Wrapper for colorcode c-function.

    This function expects a grey-scaled image in 'source' and colorcodes it
    with the given colormap 'cm'. The colorcoded map is directly written into
    source again and overwrites all data there.

    Parameters:
    -----------
    source : 1D array of doubles
        This array hold the pixels of the image that should be colorcoded.
    w : int
        The width of the image in pixels.
    h : int
        The height of the image in pixels.
    cm : [255,3] shaped int8 array
        A RGB colormap.
    colorcorrect : boolean
        Boolean that specifies if colorcorrection is needed.
    """
    return libcd.cgraphics_colorcode(source,
                                     w,
                                     h,
                                     np.float32(cm[:,0]),
                                     np.float32(cm[:,1]),
                                     np.float32(cm[:,2]),
                                     colorcorrect)

# Function for fast computation of repulsive forces between nps.
def cgraphics_np_force(item_pos_X, item_pos_Y, item_force_X, item_force_Y):
    """Wrapper for np_force c-function.

    This function computes pair-wise repulsive forces between neuron-pools.

    Parameters:
    item_pos_X : 1D double array
        This array contains x-coordinates of all neuron-pool items on the screen.
    item_pos_Y : 1D double array
        This array contains y-coordinates of all neuron-pool items on the screen.
    item_force_X : 1D double array
        In this array the resulting cummulated x-coordinate of the force is stored.
    item_force_Y : 1D double array
        In this array the resulting cummulated y-coordinate of the force is stored.
    """
    return libcd.cgraphics_np_force(item_pos_X,
                                     item_pos_Y,
                                     len(item_pos_X),
                                     item_force_X,
                                     item_force_Y)


import os
import shutil
import argparse
import numpy as np
import json
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
import torch.utils.data as data
from torch.utils.tensorboard import SummaryWriter

from sklearn.metrics import precision_recall_fscore_support
from sklearn.metrics import roc_auc_score
from sklearn.metrics import precision_recall_curve

from utils.data_utils import say, softmax
from dataset import ProcessedCNNInputDataset, ProcessedRNNInputDataset
from models.cnn import CNNMatchModel
from models.rnn import BiLSTM
from models.attn import MulInteractAttention, OneHotAttention, MLP

from utils import settings

import warnings

warnings.filterwarnings("ignore")

argparser = argparse.ArgumentParser(description="Learning to Adapt from Single Source Domain")
argparser.add_argument("--train", type=str, default="aff,author,paper,venue",
                       help="multi-source domains for training, separated with (,)")
argparser.add_argument("--test", type=str, default="aff",
                       help="target domain for testing")
argparser.add_argument("--eval_only", action="store_true")
argparser.add_argument("--batch_size", type=int, default=32)
argparser.add_argument("--max_epoch", type=int, default=300)
argparser.add_argument("--lr", type=float, default=1e-2)
argparser.add_argument("--lambda_entropy", type=float, default=0.0)
argparser.add_argument("--lambda_moe", type=float, default=1)
argparser.add_argument("--base_model", type=str, default="rnn")
argparser.add_argument("--attn-type", type=str, default="mlp")
argparser.add_argument('--train-num', default=None, help='Number of training samples')
argparser.add_argument('--n-try', type=int, default=10, help='Repeat Times')
argparser.add_argument('--n-tune', type=int, default=1, help='Number of fine-tune layers')

argparser.add_argument('--embedding-size', type=int, default=128,
                       help="Embeding size for LSTM layer")
argparser.add_argument('--hidden-size', type=int, default=32,
                       help="Hidden size for LSTM layer")
argparser.add_argument('--dropout', type=float, default=0.2,
                       help='Dropout rate (1 - keep probability).')
argparser.add_argument('--max-vocab-size', type=int, default=settings.MAX_WORD_TOKEN_NUM, help="Maximum of Vocab Size")
argparser.add_argument('--no-cuda', action='store_true', default=False, help='Disables CUDA training.')
argparser.add_argument('--matrix-size1', type=int, default=7, help='Matrix size 1.')
argparser.add_argument('--matrix-size2', type=int, default=4, help='Matrix size 2.')
argparser.add_argument('--mat1-channel1', type=int, default=4, help='Matrix1 number of channels1.')
argparser.add_argument('--mat1-kernel-size1', type=int, default=3, help='Matrix1 kernel size1.')
argparser.add_argument('--mat1-channel2', type=int, default=8, help='Matrix1 number of channel2.')
argparser.add_argument('--mat1-kernel-size2', type=int, default=2, help='Matrix1 kernel size2.')
argparser.add_argument('--mat1-hidden', type=int, default=64, help='Matrix1 hidden dim.')
argparser.add_argument('--mat2-channel1', type=int, default=4, help='Matrix2 number of channels1.')
argparser.add_argument('--mat2-kernel-size1', type=int, default=2, help='Matrix2 kernel size1.')
argparser.add_argument('--mat2-hidden', type=int, default=64, help='Matrix2 hidden dim')
argparser.add_argument('--build-index-window', type=int, default=5, help='Matrix2 hidden dim')
argparser.add_argument('--seed', type=int, default=42, help='Random seed.')
argparser.add_argument('--seed-delta', type=int, default=0, help='Random seed.')

argparser.add_argument('--weight-decay', type=float, default=1e-3,
                       help='Weight decay (L2 loss on parameters).')
argparser.add_argument('--check-point', type=int, default=2, help="Check point")
argparser.add_argument('--shuffle', action='store_true', default=True, help="Shuffle dataset")

args, _ = argparser.parse_known_args()

n_sources = len(args.train.split(','))
sources_all = ["aff", "author", "paper", "venue"]
source_to_idx = {sources_all[x]: x for x in range(len(sources_all))}
print("source to idx", source_to_idx)

def evaluate(epoch, encoders, classifier, data_loader, return_best_thrs, args, writer, thr=None):
    encoder_src, encoder_dst = encoders
    map(lambda m: m.eval(), [encoder_src, classifier, encoder_dst])
    correct = 0
    tot_cnt = 0
    y_true = []
    y_pred = []
    y_score = []
    loss = 0.

    if args.base_model == "cnn":
        for batch1, batch2, label in data_loader:
            if args.cuda:
                batch1 = batch1.cuda()
                batch2 = batch2.cuda()
                label = label.cuda()

            batch1 = Variable(batch1)
            batch2 = Variable(batch2)
            bs = len(batch1)

            _, cur_hidden = encoder_src(batch1, batch2, rsv_layer=args.n_tune)
            cur_output = classifier(cur_hidden)

            output = F.softmax(cur_output, dim=1)

            pred = output.data.max(dim=1)[1]

            loss_batch = F.nll_loss(torch.log(output), label)
            loss += bs * loss_batch.item()

            y_true += label.tolist()
            y_pred += pred.tolist()
            correct += pred.eq(label).sum()
            tot_cnt += output.size(0)
            y_score += output[:, 1].data.tolist()
    elif args.base_model == "rnn":
        for batch1, batch2, batch3, batch4, label in data_loader:
            if args.cuda:
                batch1 = batch1.cuda()
                batch2 = batch2.cuda()
                batch3 = batch3.cuda()
                batch4 = batch4.cuda()
                label = label.cuda()

            # batch1 = Variable(batch1)
            # batch2 = Variable(batch2)
            bs = len(batch1)

            _, cur_hidden = encoder_src(batch1, batch2, batch3, batch4, rsv_layer=args.n_tune)
            cur_output = classifier(cur_hidden)

            output = F.softmax(cur_output, dim=1)

            pred = output.data.max(dim=1)[1]

            loss_batch = F.nll_loss(torch.log(output), label)
            loss += bs * loss_batch.item()

            y_true += label.tolist()
            y_pred += pred.tolist()
            correct += pred.eq(label).sum()
            tot_cnt += output.size(0)
            y_score += output[:, 1].data.tolist()
    else:
        raise NotImplementedError

    if thr is not None:
        print("using threshold %.4f" % thr)
        y_score = np.array(y_score)
        y_pred = np.zeros_like(y_score)
        y_pred[y_score > thr] = 1

    loss /= tot_cnt

    prec, rec, f1, _ = precision_recall_fscore_support(y_true, y_pred, average="binary")
    auc = roc_auc_score(y_true, y_score)
    print("Loss: {:.4f}, AUC: {:.2f}, Prec: {:.2f}, Rec: {:.2f}, F1: {:.2f}".format(
        loss, auc * 100, prec * 100, rec * 100, f1 * 100))

    best_thr = None
    metric = [loss, auc, prec, rec, f1]

    if return_best_thrs:
        precs, recs, thrs = precision_recall_curve(y_true, y_score)
        f1s = 2 * precs * recs / (precs + recs)
        f1s = f1s[:-1]
        thrs = thrs[~np.isnan(f1s)]
        f1s = f1s[~np.isnan(f1s)]
        best_thr = thrs[np.argmax(f1s)]
        print("best threshold={:.4f}, f1={:.4f}".format(best_thr, np.max(f1s)))

        writer.add_scalar('val_loss',
                          loss,
                          epoch)
    else:
        writer.add_scalar('test_f1',
                          f1,
                          epoch)

    return best_thr, metric

def train_epoch(iter_cnt, encoders, classifier, train_loader_dst, args, optim_model, epoch, writer):
    encoder_src, encoder_dst = encoders

    map(lambda m: m.train(), [classifier, encoder_src, encoder_dst])
    moe_criterion = nn.NLLLoss()  # with log_softmax separated
    loss_total = 0
    n_batch = 0

    for batch in train_loader_dst:
        if args.base_model == "cnn":
            batch1, batch2, label = batch
        elif args.base_model == "rnn":
            batch1, batch2, batch3, batch4, label = batch
        else:
            raise NotImplementedError

        bs = len(label)

        iter_cnt += 1
        n_batch += 1
        if args.cuda:
            batch1 = batch1.cuda()
            batch2 = batch2.cuda()
            label = label.cuda()
            if args.base_model == "rnn":
                batch3 = batch3.cuda()
                batch4 = batch4.cuda()

        if args.base_model == "cnn":
            _, hidden_from_dst_enc = encoder_dst(batch1, batch2, rsv_layer=args.n_tune)
        elif args.base_model == "rnn":
            _, hidden_from_dst_enc = encoder_dst(batch1, batch2, batch3, batch4, rsv_layer=args.n_tune)
        else:
            raise NotImplementedError

        if args.base_model == "cnn":
            _, cur_hidden = encoder_src(batch1, batch2, rsv_layer=args.n_tune)
        elif args.base_model == "rnn":
            _, cur_hidden = encoder_src(batch1, batch2, batch3, batch4, rsv_layer=args.n_tune)
        else:
            raise NotImplementedError
        cur_output = classifier(cur_hidden)

        optim_model.zero_grad()

        output = F.softmax(cur_output, dim=1)
        loss_moe = moe_criterion(torch.log(output), label)
        loss = args.lambda_moe * loss_moe

        loss_total += loss.item()
        loss.backward()
        optim_model.step()

        if iter_cnt % 10 == 0:
            say("{} MOE loss: {:.4f}, "
                "loss: {:.4f}\n"
                .format(iter_cnt,
                        loss_moe.item(),
                        loss.data.item()
                        ))

    loss_total /= n_batch
    writer.add_scalar('training_loss',
                      loss_total,
                      epoch)

    say("\n")
    return iter_cnt

def train_single_domain_transfer(args, wf, src, repeat_seed):
    tb_dir = 'runs/{}_sup_base_{}_source_{}_train_num_{}_tune_{}_{}'.format(
        args.test, args.base_model, src, args.train_num, args.n_tune, repeat_seed)
    if os.path.exists(tb_dir) and os.path.isdir(tb_dir):
        shutil.rmtree(tb_dir)
    writer = SummaryWriter(tb_dir)

    args.cuda = not args.no_cuda and torch.cuda.is_available()
    say('cuda is available %s\n' % args.cuda)

    np.random.seed(args.seed + repeat_seed)
    torch.manual_seed(args.seed + repeat_seed)
    if args.cuda:
        torch.cuda.manual_seed(args.seed + repeat_seed)

    pretrain_emb = torch.load(os.path.join(settings.OUT_DIR, "rnn_init_word_emb.emb"))

    src_model_dir = os.path.join(settings.OUT_DIR, src)

    if args.base_model == "cnn":
        encoder_src = CNNMatchModel(input_matrix_size1=args.matrix_size1, input_matrix_size2=args.matrix_size2,
                                    mat1_channel1=args.mat1_channel1, mat1_kernel_size1=args.mat1_kernel_size1,
                                    mat1_channel2=args.mat1_channel2, mat1_kernel_size2=args.mat1_kernel_size2,
                                    mat1_hidden=args.mat1_hidden, mat2_channel1=args.mat2_channel1,
                                    mat2_kernel_size1=args.mat2_kernel_size1, mat2_hidden=args.mat2_hidden)
    elif args.base_model == "rnn":
        encoder_src = BiLSTM(pretrain_emb=pretrain_emb,
                             vocab_size=args.max_vocab_size,
                             embedding_size=args.embedding_size,
                             hidden_size=args.hidden_size,
                             dropout=args.dropout)
    else:
        raise NotImplementedError
    if args.cuda:
        encoder_src.load_state_dict(
            torch.load(os.path.join(src_model_dir,
                                    "{}-match-best-now-train-num-{}-try-{}.mdl".format(args.base_model, args.train_num,
                                                                                repeat_seed))))
    else:
        encoder_src.load_state_dict(
            torch.load(os.path.join(src_model_dir,
                                    "{}-match-best-now-train-num-{}-try-{}.mdl".format(args.base_model, args.train_num,
                                                                                repeat_seed)),
                       map_location=torch.device('cpu')))

    dst_model_dir = os.path.join(settings.OUT_DIR, args.test)
    if args.base_model == "cnn":
        encoder_dst_pretrain = CNNMatchModel(input_matrix_size1=args.matrix_size1, input_matrix_size2=args.matrix_size2,
                                             mat1_channel1=args.mat1_channel1, mat1_kernel_size1=args.mat1_kernel_size1,
                                             mat1_channel2=args.mat1_channel2, mat1_kernel_size2=args.mat1_kernel_size2,
                                             mat1_hidden=args.mat1_hidden, mat2_channel1=args.mat2_channel1,
                                             mat2_kernel_size1=args.mat2_kernel_size1, mat2_hidden=args.mat2_hidden)
    elif args.base_model == "rnn":
        encoder_dst_pretrain = BiLSTM(pretrain_emb=pretrain_emb,
                                      vocab_size=args.max_vocab_size,
                                      embedding_size=args.embedding_size,
                                      hidden_size=args.hidden_size,
                                      dropout=args.dropout)
    else:
        raise NotImplementedError

    encoder_dst_pretrain.load_state_dict(
        torch.load(os.path.join(dst_model_dir,
                                "{}-match-best-now-train-num-{}-try-{}.mdl".format(args.base_model, args.train_num,
                                                                                repeat_seed))))

    # args = argparser.parse_args()
    say(args)
    print()

    say("Transferring from %s to %s\n" % (src, args.test))

    if args.base_model == "cnn":
        train_dataset_dst = ProcessedCNNInputDataset(args.test, "train", args.train_num)
        valid_dataset = ProcessedCNNInputDataset(args.test, "valid")
        test_dataset = ProcessedCNNInputDataset(args.test, "test")

    elif args.base_model == "rnn":
        train_dataset_dst = ProcessedRNNInputDataset(args.test, "train", args.train_num)
        valid_dataset = ProcessedRNNInputDataset(args.test, "valid")
        test_dataset = ProcessedRNNInputDataset(args.test, "test")
    else:
        raise NotImplementedError

    print("train num", len(train_dataset_dst))

    train_loader_dst = data.DataLoader(
        train_dataset_dst,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=0
    )

    valid_loader = data.DataLoader(
        valid_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=0
    )

    test_loader = data.DataLoader(
        test_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=0
    )

    say("Corpus loaded.\n")

    if args.n_tune == 3:
        classifier = nn.Sequential(
            nn.Linear(encoder_src.n_out, 64),
            nn.ReLU(),
            nn.Linear(64, 16),
            nn.ReLU(),
            nn.Linear(16, 2),
        )
    elif args.n_tune == 1:
        classifier = nn.Sequential(nn.Linear(16, 2))
    # elif args.n_tune == 0:
    #     classifier = None
    else:
        raise NotImplementedError

    if args.cuda:
        encoder_dst_pretrain.cuda()
        encoder_src.cuda()
        classifier.cuda()

    requires_grad = lambda x: x.requires_grad
    task_params = list(classifier.parameters())

    if args.base_model == "cnn":
        optim_model = optim.Adagrad(
            filter(requires_grad, task_params),
            lr=args.lr,
            weight_decay=args.weight_decay
        )
    elif args.base_model == "rnn":
        optim_model = optim.Adam(
            filter(requires_grad, task_params),
            lr=args.lr,
            weight_decay=args.weight_decay
        )
    else:
        raise NotImplementedError

    say("Training will begin from scratch\n")

    iter_cnt = 0
    min_loss_val = None
    max_auc_val = None
    best_test_results = None
    weights_sources = None
    model_dir = os.path.join(settings.OUT_DIR, args.test)

    cur_src_idx = source_to_idx[src]

    for epoch in range(args.max_epoch):
        print("training epoch", epoch)

        iter_cnt = train_epoch(
            iter_cnt,
            [encoder_src, encoder_dst_pretrain], classifier,
            train_loader_dst,
            args,
            optim_model,
            epoch,
            writer
        )

        thr, metrics_val = evaluate(
            epoch,
            [encoder_src, encoder_dst_pretrain], classifier,
            valid_loader,
            True,
            args,
            writer
        )

        _, metrics_test = evaluate(
            epoch,
            [encoder_src, encoder_dst_pretrain], classifier,
            test_loader,
            False,
            args,
            writer,
            thr=thr
        )

        if min_loss_val is None or min_loss_val > metrics_val[0]:
            print("change val loss from {} to {}".format(min_loss_val, metrics_val[0]))
            min_loss_val = metrics_val[0]
            best_test_results = metrics_test
            torch.save(classifier,
                       os.path.join(model_dir, "{}_{}_classifier_from_src_{}_train_num_{}_try_{}.mdl".format(
                           args.test, args.base_model, cur_src_idx, args.train_num, repeat_seed
                       )))

    print()
    print(
        "src: {}, min valid loss {:.4f}, best test metrics: AUC: {:.2f}, Prec: {:.2f}, Rec: {:.2f}, F1: {:.2f}\n".format(
            src, min_loss_val, best_test_results[1] * 100, best_test_results[2] * 100, best_test_results[3] * 100,
                               best_test_results[4] * 100
        ))

    wf.write(
        "from src {}, min valid loss {:.4f}, best test metrics: AUC: {:.2f}, Prec: {:.2f}, Rec: {:.2f}, F1: {:.2f}\n".format(
            src, min_loss_val, best_test_results[1] * 100, best_test_results[2] * 100, best_test_results[3] * 100,
                               best_test_results[4] * 100
        ))
    writer.close()

    return min_loss_val, best_test_results[1:]

def train(args):
    model_dir = os.path.join(settings.OUT_DIR, args.test)
    wf = open(os.path.join(model_dir, "{}_{}_single_domain_trans_train_num_{}_tune_{}_results.txt".format(
        args.test, args.base_model, args.train_num, args.n_tune)), "w")
    test_idx = source_to_idx[args.test]
    for src in sources_all:
        min_loss_val_list = []
        test_results_list = []
        for t in range(args.n_try):
            cur_min_loss_val, cur_test_results = train_single_domain_transfer(args, wf, src, t)
            min_loss_val_list.append(cur_min_loss_val)
            test_results_list.append(cur_test_results)
            wf.flush()
        mean_val_loss = np.mean(min_loss_val_list)
        test_results_list = np.array(test_results_list)
        test_auc_mean = np.mean(test_results_list[:, 0])
        test_prec_mean = np.mean(test_results_list[:, 1])
        test_rec_mean = np.mean(test_results_list[:, 2])
        test_f1_mean = 2 * test_prec_mean * test_rec_mean / (test_prec_mean + test_rec_mean)
        wf.write(
            "from src {}, avg min valid loss {:.4f}, best test metrics: AUC: {:.2f}, Prec: {:.2f}, Rec: {:.2f}, F1: {:.2f}\n".format(
                src, mean_val_loss, test_auc_mean * 100, test_prec_mean * 100, test_rec_mean * 100, test_f1_mean * 100
            ))
        wf.flush()
        wf.write("\n")
    wf.write(json.dumps(vars(args)) + "\n")
    wf.close()

if __name__ == "__main__":
    train(args)
    pass


import torch
import torch.nn as nn
from attention import Attention, NewAttention
from language_model import WordEmbedding, QuestionEmbedding
from classifier import SimpleClassifier
from fc import FCNet
import torch.nn.functional as F
import torchvision as tv
import utils
import numpy as np

class vgg16_modified(nn.Module):
    def __init__(self):
        super(vgg16_modified, self).__init__()
        vgg = tv.models.vgg16(pretrained=True)
        self.vgg_features = vgg.features

    def rep_size(self):
        return 1024

    def base_size(self):
        return 512

    def forward(self,x):
        #return self.dropout2(self.relu2(self.lin2(self.dropout1(self.relu1(self.lin1(self.vgg_features(x).view(-1, 512*7*7)))))))
        features = self.vgg_features(x)

        return features

class TopDown(nn.Module):
    def __init__(self,
                 embed_hidden=300,
                 mlp_hidden=512):
        super(TopDown, self).__init__()

        self.q_emb = nn.LSTM(embed_hidden, mlp_hidden,
                             batch_first=True, bidirectional=True)
        self.q_prep = FCNet([mlp_hidden, mlp_hidden])
        self.lstm_proj = nn.Linear(mlp_hidden * 2, mlp_hidden)
        self.verb_transform = nn.Linear(embed_hidden, mlp_hidden)
        self.v_att = Attention(mlp_hidden, mlp_hidden, mlp_hidden)
        self.q_net = FCNet([mlp_hidden, mlp_hidden])
        self.v_net = FCNet([mlp_hidden, mlp_hidden])

    def forward(self, img, q):
        batch_size = img.size(0)
        w_emb = q
        lstm_out, (h, _) = self.q_emb(w_emb)
        q_emb = h.permute(1, 0, 2).contiguous().view(batch_size, -1)
        q_emb = self.lstm_proj(q_emb)

        att = self.v_att(img, q_emb)
        v_emb = (att * img).sum(1) # [batch, v_dim]

        q_repr = self.q_net(q_emb)
        v_repr = self.v_net(v_emb)
        joint_repr = q_repr * v_repr

        return joint_repr

class VerbNode(nn.Module):
    def __init__(self,
                 encoder,
                 gpu_mode,
                 verb_lookup,
                 q_word_embeddings,
                 num_verbs,
                 vqa_model,
                 mlp_hidden, embd_hidden):
        super(VerbNode, self).__init__()

        self.encoder = encoder
        self.gpu_mode = gpu_mode
        self.verb_lookup = verb_lookup
        self.q_word_embeddings = q_word_embeddings
        self.num_verbs = num_verbs
        self.vqa_model = vqa_model
        self.mlp_hidden = mlp_hidden
        self.embd_hidden = embd_hidden
        self.verb_classifier = SimpleClassifier(
            mlp_hidden, 2 * mlp_hidden, self.num_verbs, 0.5)
        self.sftmax = nn.Softmax()

    def forward(self, img, verbq, roles=None, labels=None):

        embedded_verb_q = self.q_word_embeddings(verbq)
        if roles is not None and labels is not None:
            is_id = self.encoder.question_words['is']
            is_tensor = torch.LongTensor([is_id])
            if self.gpu_mode >= 0:
                is_tensor = is_tensor.to(torch.device('cuda'))

            is_embd = self.q_word_embeddings(is_tensor)
            is_embd = is_embd.expand(roles.size(0), roles.size(1), is_embd.size(1))
            roles_labels = torch.cat((roles, is_embd, labels),1)
            embedded_verb_q = torch.cat([roles_labels, embedded_verb_q.clone()],1)

        ans_rep = self.vqa_model(img, embedded_verb_q)
        logits = self.verb_classifier(ans_rep)
        weights = self.sftmax(logits)

        weighted_verbs = torch.mm(weights, self.verb_lookup.weight)
        converted = weighted_verbs

        return logits, converted, ans_rep

class RoleNode(nn.Module):
    def __init__(self,
                 encoder,
                 gpu_mode,
                 role_lookup,
                 ans_lookup,
                 q_word_embeddings,
                 num_roles,
                 num_labels, mlp_hidden,
                 embd_hidden, vqa_model):
        super(RoleNode, self).__init__()

        self.encoder = encoder
        self.gpu_mode = gpu_mode
        self.role_lookup = role_lookup
        self.ans_lookup = ans_lookup
        self.q_word_embeddings = q_word_embeddings
        self.num_roles = num_roles
        self.num_labels = num_labels
        self.mlp_hidden = mlp_hidden
        self.embd_hidden = embd_hidden
        self.vqa_model = vqa_model
        self.v_att = Attention(mlp_hidden, mlp_hidden, mlp_hidden)
        self.context_maker = nn.Linear(self.embd_hidden*6, self.embd_hidden)
        self.contextual_verb = nn.Linear(self.embd_hidden*2 ,self.mlp_hidden)
        self.role_classifier = SimpleClassifier(
            mlp_hidden, 2 * mlp_hidden, self.num_roles + 1, 0.5)
        self.label_classifier = SimpleClassifier(
            mlp_hidden, 2 * mlp_hidden, self.num_labels + 1, 0.5)
        self.sftmax = nn.Softmax()

    def init_forward(self, img, roleq, verb, context):
        context_verb = self.contextual_verb(torch.cat((context, verb), 1))
        att = self.v_att(img, context_verb)
        v_emb = (att * img)
        v_emb_vec = v_emb.sum(1)
        role_logits = self.role_classifier(v_emb_vec)
        role_weights = self.sftmax(role_logits)
        role_ans_soft = torch.mm(role_weights, self.role_lookup.weight)
        common_roleq_phrase = self.q_word_embeddings(roleq)
        role_q = torch.cat((common_roleq_phrase, role_ans_soft.unsqueeze(1), verb.unsqueeze(1)),1)
        label_rep = self.vqa_model(v_emb, role_q)
        label_logits = self.label_classifier(label_rep)
        label_weights = self.sftmax(label_logits)
        weighted_labels = torch.mm(label_weights, self.ans_lookup.weight)
        converted = weighted_labels

        return label_logits, converted, role_logits, role_ans_soft

    def forward(self, img, roleq, verb, prev_roles, prev_labels):

        batch_size, max_role, rep_size = prev_roles.size()
        context = self.context_maker(prev_labels.view(batch_size, -1))
        context_verb = self.contextual_verb(torch.cat((context, verb), 1))
        att = self.v_att(img, context_verb)
        v_emb = (att * img)
        v_emb_vec = v_emb.sum(1)
        role_logits = self.role_classifier(v_emb_vec)
        role_weights = self.sftmax(role_logits)
        role_ans_soft = torch.mm(role_weights, self.role_lookup.weight)

        is_id = self.encoder.question_words['is']
        is_tensor = torch.LongTensor([is_id])
        if self.gpu_mode >= 0:
            is_tensor = is_tensor.to(torch.device('cuda'))

        is_embd = self.q_word_embeddings(is_tensor)
        is_embd = is_embd.expand(prev_roles.size(0), prev_roles.size(1), is_embd.size(1))

        of_tensor = torch.LongTensor([self.q_word_embeddings.weight.size(0)-1])
        if self.gpu_mode >= 0:
            of_tensor = of_tensor.to(torch.device('cuda'))

        of_embd = self.q_word_embeddings(of_tensor)
        of_embd = of_embd.expand(prev_roles.size(0), of_embd.size(0), of_embd.size(1))

        roles_labels = torch.cat((prev_roles, is_embd, prev_labels),1)

        common_roleq_phrase = self.q_word_embeddings(roleq)
        role_q = torch.cat((roles_labels, common_roleq_phrase, role_ans_soft.unsqueeze(1), of_embd, verb.unsqueeze(1)),1)

        label_rep = self.vqa_model(v_emb, role_q)
        label_logits = self.label_classifier(label_rep)
        label_weights = self.sftmax(label_logits)
        weighted_labels = torch.mm(label_weights, self.ans_lookup.weight)
        converted = weighted_labels

        return label_logits, converted, role_logits, role_ans_soft

class RecursiveGraph(nn.Module):
    def __init__(self,
                 encoder,
                 verb_lookup, role_lookup, ans_lookup,
                 q_word_embeddings,
                 num_verbs, num_roles, num_labels,
                 embd_hidden,
                 mlp_hidden,
                 max_roles,
                 gpu_mode):
        super(RecursiveGraph, self).__init__()

        self.encoder = encoder
        self.verb_lookup = verb_lookup
        self.role_lookup = role_lookup
        self.ans_lookup = ans_lookup
        self.q_word_embeddings = q_word_embeddings
        self.num_verbs = num_verbs
        self.num_roles = num_roles
        self.num_labels = num_labels
        self.emd_hidden = embd_hidden
        self.mlp_hidden = mlp_hidden
        self.max_roles = max_roles
        self.gpu_mode = gpu_mode
        self.vqa_model = TopDown()
        self.verb_node = VerbNode(self.encoder, self.gpu_mode, self.verb_lookup, self.q_word_embeddings,
                                  self.num_verbs, self.vqa_model, self.mlp_hidden, self.emd_hidden)
        self.role_node = RoleNode(self.encoder, self.gpu_mode, self.role_lookup, self.ans_lookup, self.q_word_embeddings,
                                  self.num_roles, self.num_labels, self.mlp_hidden,
                                  self.emd_hidden, self.vqa_model)

    def forward(self, img, verbq, roleq_phrase):

        verbs = []
        roles = []
        labels = []

        label_pred_all = None
        role_pred_all = None
        verb_pred_all = None

        batch_size = img.size(0)

        for iter in range(5):
            label_pred = None
            role_pred = None
            role_rep = None
            label_rep = None
            if iter == 0:
                verb_pred, verb_soft_ans, verb_ans_rep = self.verb_node(img, verbq)
                verbs.append(verb_ans_rep)
                context = torch.ones([batch_size, self.emd_hidden])

                if self.gpu_mode >= 0:
                    context = context.to(torch.device('cuda'))

                for i in range(self.max_roles):
                    label_logits, label_soft_ans, role_logits, role_ans = self.role_node.init_forward(img, roleq_phrase, verb_soft_ans, context)
                    context *= label_soft_ans

                    if i == 0:
                        label_pred = label_logits.unsqueeze(1)
                        label_rep = label_soft_ans.unsqueeze(1)
                        role_rep = role_ans.unsqueeze(1)
                        role_pred = role_logits.unsqueeze(1)
                    else:
                        label_pred = torch.cat((label_pred.clone(), label_logits.unsqueeze(1)), 1)
                        label_rep = torch.cat((label_rep.clone(), label_soft_ans.unsqueeze(1)), 1)
                        role_rep = torch.cat((role_rep.clone(), role_ans.unsqueeze(1)), 1)
                        role_pred = torch.cat((role_pred.clone(), role_logits.unsqueeze(1)), 1)

                roles.append(role_rep)
                labels.append(label_rep)

                label_pred_all = label_pred.unsqueeze(-2)
                role_pred_all = role_pred.unsqueeze(-2)
                verb_pred_all = verb_pred.unsqueeze(-2)

            else:
                verb_pred, verb_soft_ans, verb_ans_rep = self.verb_node(img, verbq, roles[-1], labels[-1])
                verbs.append(verb_ans_rep)

                mask = self.get_mask(batch_size, self.max_roles, self.emd_hidden)
                if self.gpu_mode >= 0:
                    mask = mask.to(torch.device('cuda'))

                role_prev = roles[-1]
                role_prev = role_prev.expand(mask.size(1), role_prev.size(0), role_prev.size(1), role_prev.size(2))
                role_prev = role_prev.transpose(0,1)

                label_prev = labels[-1]
                label_prev = label_prev.expand(mask.size(1), label_prev.size(0), label_prev.size(1), label_prev.size(2))
                label_prev = label_prev.transpose(0,1)

                role_prev_masked = mask * role_prev
                label_prev_masked = mask * label_prev

                #print('masked role :', role_prev_masked[0][1])

                for i in range(self.max_roles):
                    label_logits, label_soft_ans, role_logits,role_ans = self.role_node(img, roleq_phrase, verb_soft_ans,
                                                                                role_prev_masked[:,i,:,:],
                                                                                label_prev_masked[:,i,:,:])
                    if i == 0:
                        label_pred = label_logits.unsqueeze(1)
                        label_rep = label_soft_ans.unsqueeze(1)
                        role_rep = role_ans.unsqueeze(1)
                        role_pred = role_logits.unsqueeze(1)
                    else:
                        label_pred = torch.cat((label_pred.clone(), label_logits.unsqueeze(1)), 1)
                        label_rep = torch.cat((label_rep.clone(), label_soft_ans.unsqueeze(1)), 1)
                        role_rep = torch.cat((role_rep.clone(), role_ans.unsqueeze(1)), 1)
                        role_pred = torch.cat((role_pred.clone(), role_logits.unsqueeze(1)), 1)

                roles.append(role_rep)
                labels.append(label_rep)

                label_pred_all = torch.cat((label_pred_all.clone(), label_pred.unsqueeze(-2)), -2)
                role_pred_all = torch.cat((role_pred_all.clone(), role_pred.unsqueeze(-2)), -2)
                verb_pred_all = torch.cat((verb_pred_all.clone(), verb_pred.unsqueeze(-2)), -2)

        #print('sizes of pred :', label_pred_all.size(), role_pred_all.size(), verb_pred_all.size())
        label_pred_all = torch.mean(label_pred_all, -2)
        role_pred_all = torch.mean(role_pred_all, -2)
        verb_pred_all = torch.mean(verb_pred_all, -2)

        return [verb_pred_all, role_pred_all, label_pred_all ], verb_pred, label_pred

    def get_mask(self, batch_size, max_role_count, dim):
        id_mtx = torch.ones([max_role_count, max_role_count])
        np_tensor = id_mtx.numpy()
        np.fill_diagonal(np_tensor, 0)
        id_mtx = torch.from_numpy(np_tensor)
        id_mtx = (id_mtx.unsqueeze(-1)).unsqueeze(0)
        id_mtx = id_mtx.expand(batch_size, id_mtx.size(1), id_mtx.size(2), dim)

        #print('idx size :', id_mtx.size(), id_mtx[0][0], id_mtx[1][1])
        #check whether its correct

        return id_mtx

class BaseModel(nn.Module):
    def __init__(self, encoder,
                 gpu_mode,
                 embed_hidden=300,
                 mlp_hidden=512):
        super(BaseModel, self).__init__()

        self.normalize = tv.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

        self.train_transform = tv.transforms.Compose([
            tv.transforms.RandomRotation(10),
            tv.transforms.RandomResizedCrop(224),
            tv.transforms.RandomHorizontalFlip(),
            tv.transforms.ToTensor(),
            self.normalize,
        ])

        self.dev_transform = tv.transforms.Compose([
            tv.transforms.Resize(224),
            tv.transforms.CenterCrop(224),
            tv.transforms.ToTensor(),
            self.normalize,
        ])

        self.encoder = encoder
        self.gpu_mode = gpu_mode
        self.n_roles = self.encoder.get_num_roles()
        self.n_verbs = self.encoder.get_num_verbs()
        self.vocab_size = self.encoder.get_num_labels()
        self.max_role_count = self.encoder.get_max_role_count()
        self.n_role_q_vocab = len(self.encoder.question_words)

        self.conv = vgg16_modified()
        self.verb_lookup = nn.Embedding(self.n_verbs, embed_hidden)
        self.role_lookup = nn.Embedding(self.n_roles +1, embed_hidden, padding_idx=self.n_roles)
        self.ans_lookup = nn.Embedding(self.vocab_size + 1, embed_hidden, padding_idx=self.vocab_size)
        self.w_emb = nn.Embedding(self.n_role_q_vocab + 2, embed_hidden, padding_idx=self.n_role_q_vocab)

        self.vsrl_model = RecursiveGraph(self.encoder, self.verb_lookup, self.role_lookup, self.ans_lookup,
                                         self.w_emb,
                                         self.n_verbs, self.n_roles, self.vocab_size,
                                         embed_hidden,
                                         mlp_hidden,
                                         self.max_role_count,
                                         self.gpu_mode
                                         )

        self.conv_hidden = self.conv.base_size()
        self.mlp_hidden = mlp_hidden
        self.embed_hidden = embed_hidden

    def train_preprocess(self):
        return self.train_transform

    def dev_preprocess(self):
        return self.dev_transform

    def forward(self, img, verbq, roleq_phrase, verb, ans):

        img_features = self.conv(img)
        batch_size, n_channel, conv_h, conv_w = img_features.size()
        img = img_features.view(batch_size, n_channel, -1)
        img = img.permute(0, 2, 1)

        loss_rest, verb_pred, label_pred = self.vsrl_model(img, verbq, roleq_phrase)

        return loss_rest, verb_pred, label_pred

    def calculate_loss(self, verb_pred, gt_verbs, role_pred, gt_role, role_label_pred, gt_labels,args):

        batch_size = verb_pred.size()[0]
        if args.train_all:
            loss = 0
            for i in range(batch_size):
                for index in range(gt_labels.size()[1]):
                    frame_loss = 0
                    verb_loss = utils.cross_entropy_loss(verb_pred[i], gt_verbs[i])
                    #frame_loss = criterion(role_label_pred[i], gt_labels[i,index])
                    for j in range(0, self.max_role_count):
                        frame_loss += (0.2*utils.cross_entropy_loss(role_pred[i][j], gt_role[i,j] ,self.n_roles)) + \
                                      utils.cross_entropy_loss(role_label_pred[i][j], gt_labels[i,index,j] ,self.vocab_size)
                    frame_loss = (0.5*verb_loss) + frame_loss/len(self.encoder.verb2_role_dict[self.encoder.verb_list[gt_verbs[i]]])
                    #print('frame loss', frame_loss, 'verb loss', verb_loss)
                    loss += frame_loss
        else:
            #verb from pre-trained
            loss = 0
            for i in range(batch_size):
                for index in range(gt_labels.size()[1]):
                    frame_loss = 0
                    #verb_loss = utils.cross_entropy_loss(verb_pred[i], gt_verbs[i])
                    #frame_loss = criterion(role_label_pred[i], gt_labels[i,index])
                    for j in range(0, self.max_role_count):
                        frame_loss += utils.cross_entropy_loss(role_label_pred[i][j], gt_labels[i,index,j] ,self.vocab_size)
                    frame_loss = frame_loss/len(self.encoder.verb2_role_dict[self.encoder.verb_list[gt_verbs[i]]])
                    #print('frame loss', frame_loss, 'verb loss', verb_loss)
                    loss += frame_loss

        final_loss = loss/batch_size
        #print('loss :', final_loss)
        return final_loss


""" Dataset loader for the Charades-STA dataset """
import os
import csv

import h5py
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import torch.utils.data as data
import torchtext

from . import average_to_fixed_length
from core.eval import iou
from core.config import config
import datasets

class Charades(data.Dataset):
    vocab = torchtext.vocab.pretrained_aliases["glove.6B.300d"]()
    vocab.itos.extend(['<unk>'])
    vocab.stoi['<unk>'] = vocab.vectors.shape[0]
    vocab.vectors = torch.cat([vocab.vectors, torch.zeros(1, vocab.dim)], dim=0)
    word_embedding = nn.Embedding.from_pretrained(vocab.vectors)

    def __init__(self, split, rand_clip=False, training=False):
        super(Charades, self).__init__()

        self.vis_input_type = config.DATASET.VIS_INPUT_TYPE
        self.data_dir = config.DATA_DIR
        self.split = split

        self.durations = {}
        with open(os.path.join(self.data_dir, 'Charades_v1_{}.csv'.format(split))) as f:
            reader = csv.DictReader(f)
            for row in reader:
                self.durations[row['id']] = float(row['length'])

        anno_file = open(os.path.join(self.data_dir, "charades_sta_{}.txt".format(self.split)), 'r')
        annotations = []
        for line in anno_file:
            anno, sent = line.split("##")
            sent = sent.split('.\n')[0]
            vid, s_time, e_time = anno.split(" ")
            s_time = float(s_time)
            e_time = min(float(e_time), self.durations[vid])
            if s_time < e_time:
                annotations.append(
                    {'video': vid, 'times': [s_time, e_time], 'description': sent, 'duration': self.durations[vid]})
        anno_file.close()
        self.annotations = annotations
        self.rand_clip = rand_clip

    def __getitem__(self, index):
        video_id = self.annotations[index]['video']
        gt_s_time, gt_e_time = self.annotations[index]['times']
        description = self.annotations[index]['description']
        duration = self.durations[video_id]

        word_idxs = torch.tensor([self.vocab.stoi.get(w.lower(), 400000) for w in description.split()],
                                 dtype=torch.long)
        word_vectors = self.word_embedding(word_idxs)

        visual_input, visual_mask = self.get_video_features(video_id)

        if self.rand_clip:
            visual_input, duration, gt_s_time, gt_e_time = random_clip(visual_input, duration, gt_s_time, gt_e_time)

        # Time scaled to fixed size
        # visual_input = sample_to_fixed_length(visual_input, random_sampling=True)
        # visual_input = interpolate_to_fixed_length(visual_input)
        visual_input = average_to_fixed_length(visual_input)
        num_clips = config.DATASET.NUM_SAMPLE_CLIPS // config.DATASET.TARGET_STRIDE   # 16
        s_times = torch.arange(0, num_clips).float() * duration / num_clips
        e_times = torch.arange(1, num_clips + 1).float() * duration / num_clips
        overlaps = iou(torch.stack([s_times[:, None].expand(-1, num_clips),
                                    e_times[None, :].expand(num_clips, -1)], dim=2).view(-1, 2).tolist(),
                       torch.tensor([gt_s_time, gt_e_time]).tolist()).reshape(num_clips, num_clips)

        gt_s_idx = np.argmax(overlaps) // num_clips
        gt_e_idx = np.argmax(overlaps) % num_clips

        item = {
            'visual_input': visual_input,
            'anno_idx': index,
            'word_vectors': word_vectors,
            'txt_mask': torch.ones(word_vectors.shape[0], 1),
            'map_gt': torch.from_numpy(overlaps),
            'reg_gt': torch.tensor([gt_s_idx, gt_e_idx]),
            'duration': duration,
            'description': description,
        }

        return item

    def __len__(self):
        return len(self.annotations)

    def get_collate_fn(self):
        return datasets.collate_fn

    def get_video_features(self, vid):
        hdf5_file = h5py.File(os.path.join(self.data_dir, '{}_features.hdf5'.format(self.vis_input_type)), 'r')
        features = torch.from_numpy(hdf5_file[vid][:]).float()
        if config.DATASET.NORMALIZE:
            features = F.normalize(features, dim=1)
        vis_mask = torch.ones((features.shape[0], 1))
        return features, vis_mask