from __future__ import absolute_import
from __future__ import division
from os import sep

from architecture.crowd_count import CrowdCounter
import architecture.network as network
from manage_data.utils import mkdir_if_missing as mkdir_if_missing
from architecture.utils import save_results as save_results
from architecture.cam import cam as apply_cam
import numpy as np
import torch
import pandas as pd
import os.path as osp

def evaluate_model(trained_model, data_loader, model, save_test_results = False, plot_save_dir = "/tmp/", den_scale_factor = 1e3, cam=False):
    net = CrowdCounter(model = model)
    network.load_net(trained_model, net)
    net.cuda()
    net.eval()
    mae = 0.0
    mse = 0.0
    if cam:
        net.net.cam = True
    res = {'img': [], 'et': [], 'gt': [], 'abs_err': [], 'err^2': [], 'err/gt': []}

    for blob in data_loader:
        im_data = blob['data']
        gt_data = blob['gt_density']
        idx_data = blob['idx']
        im_data_norm = im_data / 127.5 - 1.
        density_map = net(im_data_norm)
        if cam in ['grad-cam', 'all']:
            apply_cam(net, blob, density_map, osp.join(plot_save_dir, 'grad-cam'), 'grad-cam')
        if cam in ['grad-cam++', 'all']:
            apply_cam(net, blob, density_map, osp.join(plot_save_dir, 'grad-cam++'), 'grad-cam++')
        #if cam in ['eigen-cam', 'all']:
        #    apply_cam(net, blob, density_map, osp.join(plot_save_dir, 'eigen-cam'), 'eigen-cam')
        if cam in ['layer-cam', 'all']:
            apply_cam(net, blob, density_map, osp.join(plot_save_dir, 'layer-cam'), 'layer-cam')
        density_map = density_map.data.cpu().numpy()
        density_map /= den_scale_factor
        gt_count = np.sum(gt_data)
        et_count = np.sum(density_map)
        mae += abs(gt_count-et_count)
        mse += ((gt_count-et_count)*(gt_count-et_count))

        res['img'].append(idx_data)
        res['et'].append(et_count)
        res['gt'].append(gt_count)
        res['abs_err'].append(abs(gt_count-et_count))
        res['err^2'].append((gt_count-et_count)*(gt_count-et_count))
        res['err/gt'].append(abs(gt_count-et_count)/gt_count)

        if save_test_results:
            print("Plotting results")
            mkdir_if_missing(plot_save_dir)
            save_results(im_data, gt_data, density_map, idx_data, plot_save_dir)
    mae = mae/data_loader.get_num_samples()
    mse = np.sqrt(mse/data_loader.get_num_samples())
    df = pd.DataFrame(res).sort_values('err/gt', axis=0, ascending=False)
    df.to_csv(osp.join(plot_save_dir, 'error_by_img.csv'), sep=',', index=False, float_format='%.5f')

    return mae,mse
