import argparse
import math
import hashlib

import yt.wrapper as yt


@yt.with_context
def format_mapper(row, context):
    gnss = row['gnss']
    loc = row['localization']

    x, y = geo_to_mercator(gnss['lon'], gnss['lat'])
    true_x, true_y = geo_to_mercator(loc['lon'], loc['lat'])

    yaw, pitch, roll = quaternion_to_euler(loc['orientation_xyzw'])
    true_heading = (rad2deg(-yaw) + 90) % 360

    yield {
        'feature_id': context.row_index,
        'source_id': hashlib.md5(row['log_time'] + row['rover']).hexdigest(),
        'image': row['image']['jpeg_data'],
        'date': int(gnss['timestamp'] * 10**9),
        'pos': {'x': x, 'y': y},
        'heading': None,
        'mds_group_id': '',
        'mds_path': '',
        'width': None,
        'height': None,
        'quality': None,
        'road_probability': None,
        'forbidden_probability': None,
        'dataset': 'selfdriving',
        'is_published': False,
        'should_be_published': False,
        'orientation': None,
        'camera_deviation': None,
        'is_hidden': None,
        'odometer_pos': None,
        'camera_orientation_rodrigues': None,

        # Additional columns
        'true_pos': {'x': true_x, 'y': true_y},
        'true_heading': true_heading,
        'mask': None,
        'intrinsics': row['image']['camera_matrix'],
    }


def quaternion_to_euler(quat):
    """Convert quaternion to euler angles"""
    x, y, z, w = quat
    sinr_cosp = 2 * (w * x + y * z)
    cosr_cosp = 1 - 2 * (x * x + y * y)
    roll = math.atan2(sinr_cosp, cosr_cosp)

    sinp = 2 * (w * y - z * x)
    if abs(sinp) >= 1:
        pitch = math.pi / 2
        pitch *= -1 if sinp < 0 else 1
    else:
        pitch = math.asin(sinp)

    siny_cosp = 2 * (w * z + x * y)
    cosy_cosp = 1 - 2 * (y * y + z * z)
    yaw = math.atan2(siny_cosp, cosy_cosp)

    return yaw, pitch, roll


def rad2deg(x):
    return x * 180 / math.pi


def geo_to_mercator(x, y):
    R = 6378137.0
    e = 0.081819190842621486448

    x = x * math.pi / 180
    y = y * math.pi / 180

    esiny = e * math.sin(y)

    tan_temp = math.tan(math.pi / 4.0 + y / 2.0)
    pow_temp = math.pow(math.tan(math.pi / 4.0 + math.asin(esiny) / 2.0), e)
    U = tan_temp / pow_temp

    return R * x, R * math.log(U)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', required=True, type=str)
    parser.add_argument('--output', required=True, type=str)
    args = parser.parse_args()

    client = yt.YtClient('hahn')

    yt.run_map(
        format_mapper,
        args.input,
        args.output,
        format=yt.YsonFormat(control_attributes_mode="iterator"),
        job_io={"control_attributes": {"enable_row_index": True}},
        client=client,
    )


if __name__ == "__main__":
    main()
