#!/usr/bin/env python

usage="""\
To use:
1. Log IMU in a T0x file (T0x_file)
2. Have RF samples logged at the same time
3. Connect a receiver to the Ettus or other sample logging system
./replayIMU.py -f imu.T04 -r 10.1.2.xx:5018
Without the "-f" we send fake static data for an Argon IMU.
Optional: use "-p admin:password" to change web username/password.
Optional: use "-o" to overwrite any existing IMU cal.  DANGER: back up first!
  After overwriting we immediately exit because the receiver reboots.
  Then you need to run without the "-o" option to send IMU data.
The specified port (e.g., 5018) must accept Trimcomm commands.
The port is only used briefly at startup.
Notes:
 - needs viewdat in path
 - tested with Titan INS
"""

import serial
import struct
from datetime import datetime
import leapseconds
import time
from collections import namedtuple
import threading
from mutils import vd2cls, get_first_str, RunCmdByLine, BinStruct
import RXTools as rx
import requests
import re
import tempfile
import os
import gzip

SECS_IN_WEEK = 24*7*60*60.

#fake_IMU_ID = 58 # uncalibrated BD940
fake_IMU_ID = 74 # calibrated BD940

class IMU_event:
    def __init__(self):
        self.IMU_ID = fake_IMU_ID
        self.gps_t_status = 5
        self.gps_week = 0
        self.gps_secs = 0.
        self.imuStatus = 1
        self.axisStatus = [3,3,3]
        self.deltaTheta = [20,18,14]
        self.deltaAccel = [-17,3,1964]
    def pack(self):
        return struct.pack('<BLBxHdLBBBxllllll',
                           0, # packet ID
                           self.IMU_ID,
                           self.gps_t_status,
                           self.gps_week,
                           self.gps_secs,
                           self.imuStatus,
                           *self.axisStatus,
                           *self.deltaTheta,
                           *self.deltaAccel)
    def full_secs(self):
        return self.gps_week*SECS_IN_WEEK + self.gps_secs

class IMU_temperature_event:
    def __init__(self):
        self.IMU_ID = fake_IMU_ID
        self.IMU_flags = 1
        self.gps_week = 0
        self.gps_secs = 0.
        self.n_rec = 3
        self.temperature=[-190,-281,-289,0,0,0]
    def pack(self):
        return struct.pack('<BLHHdBxxxllllllxxxx',
                           1, # packet ID
                           self.IMU_ID,
                           self.IMU_flags,
                           self.gps_week,
                           self.gps_secs,
                           self.n_rec,
                           *self.temperature)
    def full_secs(self):
        return self.gps_week*SECS_IN_WEEK + self.gps_secs

class ReplayIMU:
    def __init__(self, t0x_filename, all_IP_ports, user_password):
        self.conf = {}
        self.t0x_filename = t0x_filename
        self.all_IP_ports = all_IP_ports
        self.IP, self.port = all_IP_ports[0].split(':')
        self.port = int(self.port)
        self.user, self.password = user_password.split(':')
        self.web = 'http://%s:%s@%s'%(self.user,self.password,self.IP)
        self.nvc = None  # path to NVC directory: /nvc or /mnt/st/nvc

    def detect_nvc(self):
        """Detect location of receiver NVC storage. Either eCos=/nvc or Linux=/mnt/st/nvc"""
        r = requests.get( f'{self.web}/xml/dynamic/sysData.xml', timeout = 10 )
        if r.text.find('<RTOS>Linux') > 0:
            self.nvc = "/mnt/st/nvc"
        else:
            self.nvc = "/nvc"
        print(f"NVC path: {self.nvc}")

    def clear_existing_calibration( self ):
        # Clear IMU cal
        r = requests.get( self.web+'/prog/set?IMUCalibration&clear=Yes',
                          timeout=10 )
        if r.text.find('OK:') < 0:
            raise RuntimeError("Can't clear IMU calibration")
        print("Cleared IMU calibration")

    def get_default_calibration(self):
        # NVC record 48 = IMU config
        imu = BinStruct("""
        is_loaded:<B3x
        ext_imu_id:I
        is_meas_index_valid:B
        gyro_meas_index:3b
        accel_meas_index:3b
        is_temp_index_valid:B
        gyro_temperature_index:3b
        accel_temperature_index:3b
        is_scale_valid:Bx
        delta_t:d
        gyro_scale:3d
        accel_scale:3d
        is_temp_scale_bias_valid:B7x
        gyro_temperature_scale:3d
        gyro_temperature_bias:3d
        accel_temperature_scale:3d
        accel_temperature_bias:3d
        zip_cal_buffer_size:H
        zip_cal_buffer:2048B
        cal_file_name:64B
        int_imu_id:H4x
        """)
        imu.decode(bytes([0]*2304))
        # Populate fields not loaded by rec 35:5
        imu.d.is_loaded = 1
        imu.d.is_meas_index_valid = 1
        imu.d.is_temp_index_valid = 1
        imu.d.is_scale_valid = 1
        imu.d.is_temp_scale_bias_valid = 1
        return imu

    def read_rec35_sub5( self, imu, all_fields ):
        # Get config from file and update "imu"
        data = RunCmdByLine("viewdat -d35:5 --dump_avro "+self.t0x_filename)

        buf_state = -1
        if all_fields:
            while True:
                line = data.readline()
                m = re.match(r'ptr->(.*?)\.d\[(\d)\] = (.*)',line)
                if m:
                    x = getattr( imu.d, m.group(1) )
                    if type(x[0]) == float:
                        new_elem = float(m.group(3))
                    else:
                        new_elem = int(m.group(3))
                    if x[int(m.group(2))] != new_elem:
                        print("Changed ",m.groups(),x[int(m.group(2))])
                    x[int(m.group(2))] = new_elem
                    setattr( imu.d, m.group(1), x )
                elif line.startswith('ptr->delta_t'):
                    imu.d.delta_t = float(line.split()[2])
                elif line.startswith('ptr->IMU_ID'):
                    imu.d.ext_imu_id = int(line.split()[2][:-1])
                elif line.startswith('ptr->ZipCalBuffer[]:'):
                    buf_state = 0
                    break
        else:
            while True:
                line = data.readline()
                if line.startswith('ptr->ZipCalBuffer[]:'):
                    buf_state = 0
                    break

        while True:
            line = data.readline()
            if buf_state == 0:
                buf_len = int(line.split()[-1][:-1])
                buf_pos = 0
                buf_state = 1
            elif buf_state == 1:
                buf = [int(val,16) for val in line.split()]
                imu.d.zip_cal_buffer[buf_pos:buf_pos+len(buf)] = buf
                buf_pos += len(buf)
                imu.d.zip_cal_buffer_size = buf_pos
                if buf_pos >= buf_len:
                    break
        data.close()
        if buf_state != 1:
            print("WARNING - uncalibrated IMU in T0x file...")

    def overwrite_calibration(self):
        # We use calibrated IMU data, but the calibration won't match
        # what is actually on the unit.  But the IMU type and receiver
        # ID must match (or be close enough that Titan accepts it).
        self.clear_existing_calibration()
        imu = BinStruct.Empty()
        imu.d = BinStruct.Empty()
        imu.d.zip_cal_buffer = [0]*2048
        imu.d.zip_cal_buffer_size = 0
        self.read_rec35_sub5( imu, False )

        f = tempfile.NamedTemporaryFile('wb',delete=False)
        cal_data = gzip.decompress(bytearray(imu.d.zip_cal_buffer))
        f.write(cal_data)
        f.close()
        print('%s'%cal_data.decode())

        tmp_path = '/Internal/TMP2.BIN'
        r = requests.get( self.web+'/prog/delete?File&path='+tmp_path,
                          timeout=10 )
        # ignore error if TMP2.BIN does not exist

        rx.UploadFile(self.IP,self.user,self.password,
                      self.web+'/prog/Upload?DataFile&file='+tmp_path,
                      f.name,
                      response_text='OK:')
        os.remove(f.name)

        r = requests.get( self.web+'/prog/set?IMUCalibration&calfile='+tmp_path,
                          timeout=10 )
        if r.text.find('OK:') < 0:
            raise RuntimeError("Can't install IMU calibration",r.text)

        r = requests.get( self.web+'/prog/delete?File&path='+tmp_path,
                          timeout=10 )
        if r.text.find('OK:') < 0:
            raise RuntimeError("Can't delete temporary file",r.text)

        print("Uploaded new IMU calibration.  Rebooting...")

        rx.sendDColRebootRcvr( self.IP, self.port )
        time.sleep(60)

    def best_calibration_update(self):
        """Newer systems support a safer insert_fake_calibration(), so use
        that if possible.
        """
        r = requests.get( self.web+'/prog/show?FakeIMUSupport',
                          timeout=10 )
        if r.text.find('ERROR:') < 0:
            print("Creating a fake IMU config record. The original config is safe")
            self.insert_fake_calibration()
        else:
            print("Warning: overwriting IMU config!")
            self.overwrite_calibration()


    def insert_fake_calibration(self):
        # We use calibrated IMU data, but the calibration won't match
        # what is actually on the unit.  And potentially the IMU type
        # won't match either.  So.. insert a fake IMU calibration that
        # will override the existing cal. To go back to the old values,
        # the user just has to delete the fake IMU calibration file.

        imu = self.get_default_calibration()

        # Get config from file and update "imu"
        self.read_rec35_sub5( imu, True )
        imu.d.int_imu_id = 0xffff  # should be no need to set this
        # Currently in T04s we have IMU IDs:
        #  58(calibrated 59/74) ==> IMU_INT_ID_MURATA_GYRO_ACC6G(6)
        # Those IDs should work with old platforms like Argon.
        # Stella got an RTK lib update so they also work there.

        f = tempfile.NamedTemporaryFile('wb',delete=False)
        f.write(imu.encode())
        f.close()

        self.detect_nvc()
        rx.DeleteFile(self.IP,self.user,self.password,self.nvc,'91')
        rx.UploadFile(self.IP,self.user,self.password,
                      f'{self.web}/prog/Upload?DataFile&file={self.nvc}/91',
                      f.name,
                      response_text='OK:')
        os.remove(f.name)
        rx.sendDColRebootRcvr( self.IP, self.port )
        time.sleep(60)

    def init_files(self):
        lookup = [
            ('ref_to_IMU','Reference to IMU lever arm:'),
            ('ref_to_GNSS','Reference to primary GNSS lever arm:'),
            ('ref_to_GNSS_sigma','Reference to primary GNSS lever arm sigma:'),
            ('ref_to_IMU_align','Reference to IMU alignment angles:'),
            ('veh_to_ref_align','Vehicle to reference alignment angles:')]
        if self.t0x_filename is None:
            self.use_fake_data = True
            for name, search_str in lookup:
                self.conf[name] = [0.0, 0.0, 0.0]
        else:
            self.use_fake_data = False
            self.imu_data = vd2cls(self.t0x_filename,'-d35:3')
            self.temperature_data = vd2cls(self.t0x_filename,'-d35:4')
            for name, search_str in lookup:
                val = get_first_str('viewdat -d35:6 '+self.t0x_filename,
                                    search_str )
                val = val.split(b':')[-1].strip().split()
                self.conf[name] = [float(x) for x in val]


    def do_loop(self):
        self.init_files()

        need_threads = len(self.all_IP_ports) > 1
        Rcvr = namedtuple("Rcvr","IP port thread")
        self.all_rcvrs = []
        for n,IP_port in enumerate(self.all_IP_ports):
            IP,port = IP_port.split(':')
            t = None
            if need_threads:
                t = threading.Thread(target=self.run, args=(n,))
            self.all_rcvrs.append( Rcvr(IP, port, t) )
            if t is None:
                self.run(n)
            else:
                t.start()
        for rcvr in self.all_rcvrs:
            if rcvr.thread is not None:
                rcvr.thread.join()

    def run(self, n):
        rcvr = self.all_rcvrs[n]
        print("Sending DCOL startup %s:%s"%(rcvr.IP,rcvr.port))
        time.sleep(0.5)
        dcol = rx.DColSocket(rcvr.IP, int(rcvr.port))
        dcol.send_recv_ack( 0xa2, [35] )
        data = struct.pack('>Bdddd', 0x20,
                           *self.conf['ref_to_GNSS'],
                           self.conf['ref_to_GNSS_sigma'][0])
        dcol.send_recv_ack( 0xd1, data )
        data = struct.pack('>Bddd', 0x22, *self.conf['ref_to_IMU'])
        dcol.send_recv_ack( 0xd1, data )
        data = struct.pack('>Bddd', 0x23, *self.conf['ref_to_IMU_align'])
        dcol.send_recv_ack( 0xd1, data )
        data = struct.pack('>Bddd', 0x24, *self.conf['veh_to_ref_align'])
        dcol.send_recv_ack( 0xd1, data )
        dcol.send_recv_ack( 0xd1, [0x2, 1, 2] ) # enable, mapping vehicle=2
        del dcol
        time.sleep(1)

        print("Connecting to debug socket on %s"%(rcvr.IP))
        s = serial.serial_for_url('socket://%s:9999'%rcvr.IP,timeout=5)
        if self.use_fake_data:
            self.send_fake_data(n,s)
        else:
            self.send_data(n,s)
        s.close()

    def get_approx_time(self,s,delta=4.):
        print("Waiting for good time")
        while True:
            data = s.read( 10 )
            if len(data) == 10:
                break
        week, secs = struct.unpack('<Hd',data)
        # Round secs to nearest 5ms
        secs = int(secs*200)/200.0
        secs += delta
        if secs >= SECS_IN_WEEK:
            secs -= SECS_IN_WEEK
            week += 1
        print(" Got time %d %.6f"%(week,secs))
        return week, secs

    def send_fake_data(self,n,s):
        imu = IMU_event()
        temperature = IMU_temperature_event()
        week,secs = self.get_approx_time(s)
        last_time = -1.
        while True:
            if time.time() - last_time > 1.0:
                print("%d: Sending %.3f"%(n,secs))
                last_time = time.time()
            imu.gps_week = week
            imu.gps_secs = secs
            s.write(imu.pack())
            secs += 0.005
            if secs >= SECS_IN_WEEK:
                week += 1
                secs -= SECS_IN_WEEK

    def fill_imu(self,imu,i_imu):
        curr_imu = self.imu_data[i_imu]
        imu.IMU_ID = int(curr_imu.ID)
        imu.gps_week = int(curr_imu.Week)
        imu.gps_secs = curr_imu.Secs
        # no translation needed for flags:
        imu.imuStatus = int(curr_imu.Flags)
        # axisStatus needs translation:
        axisStatus = int(curr_imu.AxisStatus)
        imu.axisStatus[0] = 0
        imu.axisStatus[1] = 0
        imu.axisStatus[2] = 0
        if( axisStatus & 0x01 ): imu.axisStatus[0] |= 1 # gyro X present
        if( axisStatus & 0x02 ): imu.axisStatus[1] |= 1 # gyro Y present
        if( axisStatus & 0x04 ): imu.axisStatus[2] |= 1 # gyro Z present
        if( axisStatus & 0x08 ): imu.axisStatus[0] |= 2 # accel X present
        if( axisStatus & 0x10 ): imu.axisStatus[1] |= 2 # accel Y present
        if( axisStatus & 0x20 ): imu.axisStatus[2] |= 2 # accel Z present
        # old firmware may not have higher flags...
        imu.deltaTheta[0] = int(curr_imu.GyroX)
        imu.deltaTheta[1] = int(curr_imu.GyroY)
        imu.deltaTheta[2] = int(curr_imu.GyroZ)
        imu.deltaAccel[0] = int(curr_imu.AccX)
        imu.deltaAccel[1] = int(curr_imu.AccY)
        imu.deltaAccel[2] = int(curr_imu.AccZ)
        return i_imu+1

    def fill_temperature(self,T,i_T):
        T.n_rec = 0
        while True:
            curr_T = self.temperature_data[i_T]
            T.IMU_ID = int(curr_T.ID)
            T.gps_week = int(curr_T.Week)
            T.gps_secs = curr_T.Secs
            # T.Flags = assume data is OK for now
            T.temperature[ T.n_rec ] = int(curr_T.Temperature)
            T.n_rec += 1
            i_T += 1
            if abs(curr_T.Secs - self.temperature_data[i_T].Secs) > 1e-3:
                break
        return i_T

    def get_approx_imu_start(self,week,secs):
        i_imu = 0
        start_t = week*SECS_IN_WEEK + secs
        while True:
            if i_imu >= len(self.imu_data):
                raise RuntimeError('get_approx_imu_start error',week,secs,self.imu_data[-1].Week,self.imu_data[-1].Secs)
            file_t = (self.imu_data[i_imu].Week*SECS_IN_WEEK
                      + self.imu_data[i_imu].Secs)
            if file_t >= start_t:
                break
            i_imu += 1
        return i_imu

    def get_approx_T_start(self,week,secs):
        i_T = 0
        start_t = week*SECS_IN_WEEK + secs
        while True:
            file_t = (self.temperature_data[i_T].Week*SECS_IN_WEEK
                      + self.temperature_data[i_T].Secs)
            if file_t >= start_t:
                break
            i_T += 1
        return i_T

    def send_data(self,n,s):
        imu = IMU_event()
        T = IMU_temperature_event()
        week,secs = self.get_approx_time(s)
        i_imu = self.get_approx_imu_start(week,secs)
        i_T = self.get_approx_T_start(week,secs)
        i_imu = self.fill_imu( imu, i_imu )
        i_T = self.fill_temperature( T, i_T )
        last_time = -1.
        while True:
            if time.time() - last_time > 1.0:
                print("%d: Sending IMU %.3f temperature %.3f"% (
                      n,
                      imu.gps_secs,
                      T.gps_secs))
                last_time = time.time()

            # Send oldest data first
            if imu.full_secs() < T.full_secs():
                s.write(imu.pack())
                i_imu = self.fill_imu( imu, i_imu )
            else:
                s.write(T.pack())
                i_T = self.fill_temperature( T, i_T )


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description=usage,
                         formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('-f',"--T0x_file", help="File with IMU data.  If not specified, then we'll transmit fake static data.")
    parser.add_argument('-r',"--rcvrs", nargs='+', help="List of receiver IP:ports", required=True)
    parser.add_argument('-p',"--password", help="web User:Password", default="admin:password")
    parser.add_argument('-o',"--overwrite", help="DANGEROUS for old systems: overwrite any existing IMU calibration, so back up first! Safe for new systems that have prog iface Show?FakeIMUSupport command.", action="store_true")
    args = parser.parse_args()
    r = ReplayIMU(args.T0x_file, args.rcvrs, args.password)
    if args.overwrite:
        r.best_calibration_update()
    else:
        r.do_loop()