#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Apr  6 14:35:43 2020

Create Histology Session on IBL Alyx

This script uploads a Histology Session to IBL Alyx including sample ID,
sample Reception and Imaging dates, and the affine matrix. The tilt (about
ML dimension), yaw (about DV dimension), roll (about AP dimension), dv_scale,
ap_scale, ml_scale values are computed from the affine matrix.


Parameters
----------
id : str
    The subject id
sample_imaging_date : str
    Date of Sample Imaging, given in format DD-MM-YYYY
sample_reception_date : str
    Date of Sample Reception, given in format DD-MM-YYYY
affine_matrix : str
    4x4 Affine Transformation Matrix mapping sample2ARA/, given in format
    "0.912838 0.038861 0.142005 -0.054963 0.820837 0.034729 -0.135135
    -0.065278 0.929171 -18.357938 -19.875191 -5.734609"


Returns
-------
none
"""
# Author: Steven West (main), Olivier Winter, Gaelle Chapuis

import sys
import math
import datetime
import json
from json import JSONEncoder

import numpy as np
from one.api import ONE

import ibllib.time


# override default method of JSONEncoder to implement custom NumPy JSON serialization.
# see https://pynative.com/python-serialize-numpy-ndarray-into-json/
class NumpyArrayEncoder(JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return JSONEncoder.default(self, obj)


# collect ARGS
subject = sys.argv[1]  # subject id is in the first ARG
sample_imaging_date = sys.argv[2]
sample_reception_date = sys.argv[3]
affine_matrix = sys.argv[4]


# check variables are valid:
if len(sample_imaging_date) != 10 \
    or sample_imaging_date[2:3] != "-" \
        or sample_imaging_date[5:6] != "-":
    print("sample_imaging_date not in correct format")
    sys.exit()

if len(sample_reception_date) != 10 \
    or sample_reception_date[2:3] != "-" \
        or sample_reception_date[5:6] != "-":
    print("sample_reception_date not in correct format")
    sys.exit()

if affine_matrix.count(" ") != 11:
    print("affine_matrix not in correct format")
    sys.exit()


# Break date of imaging and reception into D-M-Y:
imD = int(sample_imaging_date[0:2])
imM = int(sample_imaging_date[3:5])
imY = int(sample_imaging_date[6:10])

reD = int(sample_reception_date[0:2])
reM = int(sample_reception_date[3:5])
reY = int(sample_reception_date[6:10])

sample_imaging_date = datetime.date(imY, imM, imD)  # Format: y - m - d
sample_reception_date = datetime.date(reY, reM, reD)  # Format: y - m - d


# Create 4x4 affine_matrix in Homogenous Coordinates:
aff_list = [float(x) for x in affine_matrix.split()]

affine_matrix = np.zeros((4, 4))
affine_matrix[3, 3] = 1.0

# fill 3x3 affine matrix in row-major ordering!
affine_matrix[0, 0] = aff_list[0]
affine_matrix[0, 1] = aff_list[1]
affine_matrix[0, 2] = aff_list[2]

affine_matrix[1, 0] = aff_list[3]
affine_matrix[1, 1] = aff_list[4]
affine_matrix[1, 2] = aff_list[5]

affine_matrix[2, 0] = aff_list[6]
affine_matrix[2, 1] = aff_list[7]
affine_matrix[2, 2] = aff_list[8]

# the translations fill in the last column of the matrix:
affine_matrix[0, 3] = aff_list[9]
affine_matrix[1, 3] = aff_list[10]
affine_matrix[2, 3] = aff_list[11]


# Extract the scale and rotation matrix:
sca_rot_mat = np.zeros((3, 3))

# fill 3x3 affine matrix in row-major ordering!
sca_rot_mat[0, 0] = aff_list[0]
sca_rot_mat[0, 1] = aff_list[1]
sca_rot_mat[0, 2] = aff_list[2]

sca_rot_mat[1, 0] = aff_list[3]
sca_rot_mat[1, 1] = aff_list[4]
sca_rot_mat[1, 2] = aff_list[5]

sca_rot_mat[2, 0] = aff_list[6]
sca_rot_mat[2, 1] = aff_list[7]
sca_rot_mat[2, 2] = aff_list[8]

# extract scale:
sx = math.sqrt(((sca_rot_mat[0, 0]**2) + (sca_rot_mat[1, 0]**2) + (sca_rot_mat[2, 0]**2)))
sy = math.sqrt(((sca_rot_mat[0, 1]**2) + (sca_rot_mat[1, 1]**2) + (sca_rot_mat[2, 1]**2)))
sz = math.sqrt(((sca_rot_mat[0, 2]**2) + (sca_rot_mat[1, 2]**2) + (sca_rot_mat[2, 2]**2)))


# extract composite rotation matrix:
rot_mat = np.zeros((3, 3))

rot_mat[0, 0] = sca_rot_mat[0, 0] / sx
rot_mat[1, 0] = sca_rot_mat[1, 0] / sx
rot_mat[2, 0] = sca_rot_mat[2, 0] / sx

rot_mat[0, 1] = sca_rot_mat[0, 1] / sy
rot_mat[1, 1] = sca_rot_mat[1, 1] / sy
rot_mat[2, 1] = sca_rot_mat[2, 1] / sy

rot_mat[0, 2] = sca_rot_mat[0, 2] / sz
rot_mat[1, 2] = sca_rot_mat[1, 2] / sz
rot_mat[2, 2] = sca_rot_mat[2, 2] / sz


# extract individual rotations in x y z:
rx = math.atan2(rot_mat[2, 1], rot_mat[2, 2])
ry = math.atan2((-rot_mat[2, 0]), math.sqrt(((rot_mat[2, 1]**2) + (rot_mat[2, 2]**2))))
rz = math.atan2(rot_mat[1, 0], rot_mat[0, 0])


# Upload to IBL alyx:
one = ONE(base_url='https://dev.alyx.internationalbrainlab.org')


TASK_PROTOCOL = 'SWC_Histology_Serial2P_v0.0.1'


json_note = {
    'sample_reception_date': ibllib.time.date2isostr(sample_reception_date),
    'elastix_affine_transform': affine_matrix,
    'tilt': rx,
    'yaw': ry,
    'roll': rz,
    'dv_scale': sy,
    'ap_scale': sz,
    'ml_scale': sx
}

# use dump() to properly encode np array:
json_note = json.dumps(json_note, cls=NumpyArrayEncoder)

ses_ = {
    'subject': subject,
    'users': ['steven.west'],
    'location': 'serial2P_01',
    'procedures': ['Histology'],
    'lab': 'mrsicflogellab',
    # 'project': project['name'],
    # 'type': 'Experiment',
    'task_protocol': TASK_PROTOCOL,
    'number': 1,
    'start_time': ibllib.time.date2isostr(sample_imaging_date),  # Saving only the date
    # 'end_time': ibllib.time.date2isostr(end_time) if end_time else None,
    # 'n_correct_trials': n_correct_trials,
    # 'n_trials': n_trials,
    'json': json_note
}

# overwrites the session if it already exists
ses_date = ibllib.time.date2isostr(sample_imaging_date)[:10]
ses = one.alyx.rest('sessions', 'list', subject=subject, number=1,
                    date_range=[ses_date, ses_date], no_cache=True)
if len(ses) > 0:
    one.alyx.rest('sessions', 'delete', ses[0]['url'])

session = one.alyx.rest('sessions', 'create', data=ses_)
