#include "vna_litevna.h"

/*
 LiteVNA 64 - protocollo USB SAA2
 --------------------------------
 Registro 0x00: sweepStartHz, uint64 little endian
 Registro 0x10: sweepStepHz,  uint64 little endian
 Registro 0x20: sweepPoints,  uint16 little endian
 FIFO     0x30: record da 32 byte

 Correzione implementata:
 - S11: calibrazione SOL a 3 termini (directivity, source match,
        reflection tracking), standard ideali OPEN/SHORT/LOAD.
 - S21: normalizzazione THRU con sottrazione ISOLATION.

 Non e' una calibrazione 12-term completa.
*/

#define VNA_OP_INDICATE  0x0d
#define VNA_OP_READFIFO  0x18
#define VNA_OP_WRITE     0x20
#define VNA_OP_WRITE2    0x21
#define VNA_OP_WRITE8    0x23

#define VNA_REG_START    0x00
#define VNA_REG_STEP     0x10
#define VNA_REG_POINTS   0x20
#define VNA_REG_FIFO     0x30
#define VNA_REG_VARIANT  0xf0
#define VNA_REG_PROTO    0xf1
#define VNA_REG_AVERAGE  0x40

#define VNA_FIFO_RECORD_SIZE 32
#define VNA_FIFO_CHUNK       8
#define VNA_IO_TIMEOUT       500
#define VNA_SCAN_TIMEOUT     30000




/* Standard di calibrazione acquisiti */
static VNA_COMPLEX cal_open[VNA_MAX_POINTS];
static VNA_COMPLEX cal_short[VNA_MAX_POINTS];
static VNA_COMPLEX cal_load[VNA_MAX_POINTS];
static VNA_COMPLEX cal_isol[VNA_MAX_POINTS];
static VNA_COMPLEX cal_thru[VNA_MAX_POINTS];

/* Coefficienti */
static VNA_COMPLEX cal_ed[VNA_MAX_POINTS];
static VNA_COMPLEX cal_es[VNA_MAX_POINTS];
static VNA_COMPLEX cal_er[VNA_MAX_POINTS];


void vna_clear_calibration(void)
{
    have_open=0;
    have_short=0;
    have_load=0;
    have_isol=0;
    have_thru=0;
    vna_cal_valid=0;

    memset(cal_open,0,sizeof(cal_open));
    memset(cal_short,0,sizeof(cal_short));
    memset(cal_load,0,sizeof(cal_load));
    memset(cal_isol,0,sizeof(cal_isol));
    memset(cal_thru,0,sizeof(cal_thru));
    memset(cal_ed,0,sizeof(cal_ed));
    memset(cal_es,0,sizeof(cal_es));
    memset(cal_er,0,sizeof(cal_er));
}


/* ---------------------------------------------------------- */
static VNA_COMPLEX cx(float re,float im)
{
    VNA_COMPLEX z;
    z.re=re;
    z.im=im;
    return z;
}

static VNA_COMPLEX cx_add(VNA_COMPLEX a,VNA_COMPLEX b)
{
    return cx(a.re+b.re,a.im+b.im);
}

static VNA_COMPLEX cx_sub(VNA_COMPLEX a,VNA_COMPLEX b)
{
    return cx(a.re-b.re,a.im-b.im);
}

static VNA_COMPLEX cx_mul(VNA_COMPLEX a,VNA_COMPLEX b)
{
    return cx(a.re*b.re-a.im*b.im,a.re*b.im+a.im*b.re);
}

static VNA_COMPLEX cx_div(VNA_COMPLEX a,VNA_COMPLEX b)
{
    double den=(double)b.re*b.re+(double)b.im*b.im;
    if (den < 1.0e-30) return cx(0,0);
    return cx((float)(((double)a.re*b.re+(double)a.im*b.im)/den),
    (float)(((double)a.im*b.re-(double)a.re*b.im)/den));
}

/* ---------------------------------------------------------- */
static U16 get_u16(const U8 *p)
{
    return (U16)((U16)p[0] | ((U16)p[1]<<8));
}

static S32 get_s32(const U8 *p)
{
    U32 u=(U32)p[0] |
    ((U32)p[1]<<8) |
    ((U32)p[2]<<16) |
    ((U32)p[3]<<24);
    return (S32)u;
}

static void put_u16(U8 *p,U16 v)
{
    p[0]=(U8)v;
    p[1]=(U8)(v>>8);
}

static void put_u64(U8 *p,unsigned long long v)
{
    int i;
    for(i=0;i<8;i++)
    {
        p[i]=(U8)v;
        v>>=8;
    }
}

/* ---------------------------------------------------------- */
void vna_update_port_name(void)
{
#ifdef MSWINDOWS
    sprintf(vna_port_name,"COM%d",(int)vna_port_index+1);
#else
    sprintf(vna_port_name,"/dev/ttyACM%d",(int)vna_port_index);
#endif
}

void vna_post_command(int command)
{
    if (command==VNA_CMD_ABORT)
    {
        vna_abort=1;
        return;
    }

    if (vna_command!=VNA_CMD_NONE || vna_state==VNA_STATE_BUSY)
    {
        strcpy(vna_status,"Command ignored: VNA busy");
        return;
    }
    vna_command=command;
}

/* ---------------------------------------------------------- */
static void vna_flush_rx(void)
{
    while(!com_rx_empty(vna_com)) com_rx(vna_com);
}

static int vna_write_bytes(const U8 *src,int len)
{
    int n;
    if (!vna_com->open) return 0;
    n=com_txsl(vna_com,(unsigned char *)src,len);
    return (n>=0);
}

static int vna_read_exact(U8 *dst,int len,int timeout)
{
    int n=0,c;
    U32 start=(U32)_time(0);

    while(n<len)
    {
        while(!com_rx_empty(vna_com))
        {
            c=com_rx(vna_com);
            if(c>=0) dst[n++]=(U8)c;
            if(n>=len) return 1;
        }

        if(((U32)_time(0)-start)>(U32)timeout) return 0;
        wait(1);
    }
    return 1;
}

static int vna_indicate(void)
{
    U8 tx=VNA_OP_INDICATE;
    U8 rx=0;
    vna_flush_rx();
    if(!vna_write_bytes(&tx,1)) return 0;
    if(!vna_read_exact(&rx,1,VNA_IO_TIMEOUT)) return 0;
    return rx==0x32;
}

static int vna_write1(U8 addr,U8 value)
{
    U8 b[3]={VNA_OP_WRITE,addr,value};
    return vna_write_bytes(b,3);
}

static int vna_write2(U8 addr,U16 value)
{
    U8 b[4];
    b[0]=VNA_OP_WRITE2;
    b[1]=addr;
    put_u16(&b[2],value);
    return vna_write_bytes(b,4);
}

static int vna_write8(U8 addr,unsigned long long value)
{
    U8 b[10];
    b[0]=VNA_OP_WRITE8;
    b[1]=addr;
    put_u64(&b[2],value);
    return vna_write_bytes(b,10);
}

static int vna_read_fifo(U8 records,U8 *buffer)
{
    U8 b[3]={VNA_OP_READFIFO,VNA_REG_FIFO,records};
    if(!vna_write_bytes(b,3)) return 0;
    return vna_read_exact(buffer,(int)records*VNA_FIFO_RECORD_SIZE,VNA_IO_TIMEOUT);
}

/* ---------------------------------------------------------- */
static int vna_connect_device(void)
{
    int ret;

    strcpy(vna_status,"Opening port...");
    sprintf(vna_port_name,"%s",Co1->items[-1]);

    ret=assign_com(vna_com,vna_port_name);
    if(ret<0)
    {
        strcpy(vna_status,"assign_com error");
        return 0;
    }

    ret=com_open(vna_com,115200);
    if(ret<0)
    {
        strcpy(vna_status,"Port open error");
        com_close(vna_com);
        return 0;
    }

    protocol_mode(vna_com,0);
    vna_flush_rx();

    if(!vna_indicate())
    {
        com_close(vna_com);
        strcpy(vna_status,"LiteVNA not responding");
        return 0;
    }

    strcpy(vna_status,"LiteVNA connected");
    vna_state=VNA_STATE_CONNECTED;
    return 1;
}

static void vna_disconnect_device(void)
{
    if(vna_com->open) com_close(vna_com);
    vna_state=VNA_STATE_IDLE;
    strcpy(vna_status,"Not connected");
}

/* ---------------------------------------------------------- */
static int vna_validate_sweep(void)
{
    if(vna_points<2 || vna_points>VNA_MAX_POINTS)
    {
        strcpy(vna_status,"Invalid number of points");
        return 0;
    }
    if(vna_f_stop<=vna_f_start)
    {
        strcpy(vna_status,"Invalid frequencies");
        return 0;
    }
    return 1;
}

static int vna_set_sweep(void)
{
    unsigned long long start=(unsigned long long)vna_f_start;
    unsigned long long stop =(unsigned long long)vna_f_stop;
    unsigned long long step =(stop-start)/(unsigned long long)(vna_points-1);
   
    if(!vna_write8(VNA_REG_START,start)) return 0;
    if(!vna_write8(VNA_REG_STEP,step)) return 0;
    if(!vna_write2(VNA_REG_POINTS,vna_points)) return 0;
    if (slow)  if(!vna_write1(VNA_REG_AVERAGE,4)) return 0;
    if (!slow) if(!vna_write1(VNA_REG_AVERAGE,1)) return 0;
/*  Scrivendo qualsiasi valore a 0x30 si svuota la FIFO. */
    if(!vna_write1(VNA_REG_FIFO,0)) return 0;
    return 1;
}

static void vna_decode_record(const U8 *r)
{
    S32 fwd_re=get_s32(r+0x00);
    S32 fwd_im=get_s32(r+0x04);
    S32 s11_re=get_s32(r+0x08);
    S32 s11_im=get_s32(r+0x0c);
    S32 s21_re=get_s32(r+0x10);
    S32 s21_im=get_s32(r+0x14);
    U16 index=get_u16(r+0x18);

    VNA_COMPLEX ref=cx((float)fwd_re,(float)fwd_im);
    VNA_COMPLEX a=cx_div(cx((float)s11_re,(float)s11_im),ref);
    VNA_COMPLEX b=cx_div(cx((float)s21_re,(float)s21_im),ref);

    if(index>=vna_points || index>=VNA_MAX_POINTS) return;

    vna_frequency[index]=vna_f_start+
    (U32)(((unsigned long long)(vna_f_stop-vna_f_start)*index)/
    (unsigned long long)(vna_points-1));

    vna_s11_raw_re[index]=a.re;
    vna_s11_raw_im[index]=a.im;
    vna_s21_raw_re[index]=b.re;
    vna_s21_raw_im[index]=b.im;
}

/* ---------------------------------------------------------- */
static void vna_apply_calibration(void)
{
    int i;
    for(i=0;i<vna_points;i++)
    {
        VNA_COMPLEX m11=cx(vna_s11_raw_re[i],vna_s11_raw_im[i]);
        VNA_COMPLEX m21=cx(vna_s21_raw_re[i],vna_s21_raw_im[i]);
        VNA_COMPLEX s11=m11;
        VNA_COMPLEX s21=m21;

        if(vna_cal_valid)
        {
            VNA_COMPLEX d=cx_sub(m11,cal_ed[i]);
            s11=cx_div(d,cx_add(cal_er[i],cx_mul(cal_es[i],d)));
            s21=cx_div(cx_sub(m21,cal_isol[i]),
            cx_sub(cal_thru[i],cal_isol[i]));
        }

        vna_s11_re[i]=s11.re;
        vna_s11_im[i]=s11.im;
        vna_s21_re[i]=s21.re;
        vna_s21_im[i]=s21.im;
    }
}

static int vna_scan_raw(void)
{
    U8 buffer[VNA_FIFO_CHUNK*VNA_FIFO_RECORD_SIZE];
    U8 seen[VNA_MAX_POINTS];
    int count=0;
    int i,j;
    U32 start_time;

    if(vna_state!=VNA_STATE_CONNECTED && vna_state!=VNA_STATE_BUSY)
    {
        strcpy(vna_status,"VNA not connected");
        return 0;
    }
    if(!vna_validate_sweep()) return 0;

    memset(seen,0,sizeof(seen));
    vna_abort=0;
    vna_acquired=0;
    vna_data_valid=0;

    if(!vna_set_sweep())
    {
        strcpy(vna_status,"Sweep setup error");
        return 0;
    }

    start_time=(U32)_time(0);

    while(count<vna_points)
    {
        int request=VNA_FIFO_CHUNK;
        if(vna_abort)
        {
            strcpy(vna_status,"Scan aborted");
            return 0;
        }
        if(((U32)_time(0)-start_time)>VNA_SCAN_TIMEOUT)
        {
            strcpy(vna_status,"Scan timeout");
            return 0;
        }

        if(!vna_read_fifo((U8)request,buffer))
        {
            strcpy(vna_status,"FIFO read error");
            return 0;
        }

        for(j=0;j<request;j++)
        {
            const U8 *r=&buffer[j*VNA_FIFO_RECORD_SIZE];
            U16 index=get_u16(r+0x18);
            vna_decode_record(r);
            if(index<vna_points && !seen[index])
            {
                seen[index]=1;
                count++;
                vna_acquired=count;
            }
        }
    }
    start_time=(U32)_time(0)-start_time;
    MD777=start_time;
    vna_apply_calibration();
    vna_data_valid=1;
    return 1;
}

/* ---------------------------------------------------------- */
static void vna_copy_standard(VNA_COMPLEX *dst,int channel)
{
    int i;
    for(i=0;i<vna_points;i++)
    {
        if(channel==11)
        dst[i]=cx(vna_s11_raw_re[i],vna_s11_raw_im[i]);
        else
        dst[i]=cx(vna_s21_raw_re[i],vna_s21_raw_im[i]);
    }
}

static int vna_compute_calibration(void)
{
    int i;

    if(!have_open || !have_short || !have_load ||
    !have_isol || !have_thru)
    {
        strcpy(vna_status,"Incomplete calibration");
        return 0;
    }

    for(i=0;i<vna_points;i++)
    {
        VNA_COMPLEX ed=cal_load[i];
        VNA_COMPLEX a=cx_sub(cal_open[i],ed);
        VNA_COMPLEX b=cx_sub(cal_short[i],ed);
        VNA_COMPLEX es=cx_div(cx_add(a,b),cx_sub(a,b));
        VNA_COMPLEX er=cx_mul(a,cx_sub(cx(1,0),es));

        cal_ed[i]=ed;
        cal_es[i]=es;
        cal_er[i]=er;
    }

    vna_cal_valid=1;
    strcpy(vna_status,"Calibration active");
    return 1;
}

/* ---------------------------------------------------------- */
#define CAL_MAGIC 0x564e4143UL
#define CAL_VERSION 1

typedef struct CAL_HEADER
{
    U32 magic;
    U32 version;
    U32 start;
    U32 stop;
    U16 points;
    U16 reserved;
} CAL_HEADER;

static int write_complex(FILE *f,VNA_COMPLEX *v,int n)
{
    return fwrite(v,sizeof(VNA_COMPLEX),n,f)==(U32)n;
}

static int read_complex(FILE *f,VNA_COMPLEX *v,int n)
{
    return fread(v,sizeof(VNA_COMPLEX),n,f)==(U32)n;
}

static int vna_save_calibration(void)
{
    FILE *f;
    CAL_HEADER h;

    if(!vna_cal_valid)
    {
        strcpy(vna_status,"No calibration to save");
        return 0;
    }

    f=fopen(vna_cal_file,"wb");
    if(!f)
    {
        strcpy(vna_status,"Calibration file open error");
        return 0;
    }

    h.magic=CAL_MAGIC;
    h.version=CAL_VERSION;
    h.start=vna_f_start;
    h.stop=vna_f_stop;
    h.points=vna_points;
    h.reserved=0;

    if(fwrite(&h,sizeof(h),1,f)!=1 ||
        !write_complex(f,cal_open,vna_points) ||
        !write_complex(f,cal_short,vna_points) ||
        !write_complex(f,cal_load,vna_points) ||
        !write_complex(f,cal_isol,vna_points) ||
    !write_complex(f,cal_thru,vna_points))
    {
        fclose(f);
        strcpy(vna_status,"Calibration write error");
        return 0;
    }

    fclose(f);
    strcpy(vna_status,"Calibration saved");
    return 1;
}

static int vna_load_calibration(char *FNAME)
{
    FILE *f;
    CAL_HEADER h;

    f=fopen(vna_cal_file,"rb");
    if(!f)
    {
        strcpy(vna_status,"Calibration file not found");
        return 0;
    }

    if(fread(&h,sizeof(h),1,f)!=1 ||
        h.magic!=CAL_MAGIC ||
        h.version!=CAL_VERSION ||
        h.points<2 ||
    h.points>VNA_MAX_POINTS)
    {
        fclose(f);
        strcpy(vna_status,"Invalid calibration file");
        return 0;
    }

    vna_f_start=h.start;
    vna_f_stop=h.stop;
    vna_points=h.points;
    f_start=vna_f_start/1000000.;
    f_stop=vna_f_stop/1000000.;

    if(!read_complex(f,cal_open,vna_points) ||
        !read_complex(f,cal_short,vna_points) ||
        !read_complex(f,cal_load,vna_points) ||
        !read_complex(f,cal_isol,vna_points) ||
    !read_complex(f,cal_thru,vna_points))
    {
        fclose(f);
        strcpy(vna_status,"Incomplete calibration file");
        return 0;
    }

    fclose(f);
    have_open=have_short=have_load=have_isol=have_thru=1;
    return vna_compute_calibration();
}

/* ---------------------------------------------------------- */
static void vna_capture(int command)
{
    const char *name="";

    if(command==VNA_CMD_CAL_OPEN)
    {
        vna_clear_calibration();
        name="OPEN";
    }
    if(command==VNA_CMD_CAL_SHORT) name="SHORT";
    if(command==VNA_CMD_CAL_LOAD)  name="LOAD";
    if(command==VNA_CMD_CAL_ISOL)  name="ISOLATION";
    if(command==VNA_CMD_CAL_THRU)  name="THRU";

    sprintf(vna_status,"Acquisition %s...",name);

    if(!vna_scan_raw()) return;

    if(command==VNA_CMD_CAL_OPEN)
    {
        vna_copy_standard(cal_open,11);
        have_open=1;
    }
    else if(command==VNA_CMD_CAL_SHORT)
    {
        vna_copy_standard(cal_short,11);
        have_short=1;
    }
    else if(command==VNA_CMD_CAL_LOAD)
    {
        vna_copy_standard(cal_load,11);
        have_load=1;
    }
    else if(command==VNA_CMD_CAL_ISOL)
    {
        vna_copy_standard(cal_isol,21);
        have_isol=1;
    }
    else if(command==VNA_CMD_CAL_THRU)
    {
        vna_copy_standard(cal_thru,21);
        have_thru=1;
    }

    sprintf(vna_status,"%s acquised",name);
}

/* ---------------------------------------------------------- */
void vna_task(void)
{
    int cmd;

    for(;;)
    {
        cmd=vna_command;
        if(cmd==VNA_CMD_NONE)
        {
            wait(1);
            continue;
        }

        vna_command=VNA_CMD_NONE;

        if(cmd==VNA_CMD_ABORT)
        {
            vna_abort=1;
            continue;
        }

        if(cmd==VNA_CMD_TERMINATE)
        {
            vna_disconnect_device();
            return;
        }

        if(cmd==VNA_CMD_CONNECT)
        {
            if(vna_state==VNA_STATE_IDLE || vna_state==VNA_STATE_ERROR)
            {
                vna_state=VNA_STATE_BUSY;
                if(!vna_connect_device()) vna_state=VNA_STATE_ERROR;
            }
            continue;
        }

        if(cmd==VNA_CMD_DISCONNECT)
        {
            vna_disconnect_device();
            continue;
        }

        if(cmd==VNA_CMD_CAL_RECALL)
        {
            vna_state=VNA_STATE_BUSY;
            if(!vna_load_calibration(vna_cal_file)) vna_state=VNA_STATE_ERROR;
            else if(vna_com->open) vna_state=VNA_STATE_CONNECTED;
            else vna_state=VNA_STATE_IDLE;
            continue;
        }

        if(cmd==VNA_CMD_CAL_SAVE)
        {
            vna_state=VNA_STATE_BUSY;
            vna_save_calibration();
            if(vna_com->open) vna_state=VNA_STATE_CONNECTED;
            else vna_state=VNA_STATE_IDLE;
            continue;
        }

        if(vna_state!=VNA_STATE_CONNECTED)
        {
            strcpy(vna_status,"Connect LiteVNA first");
            continue;
        }

        vna_state=VNA_STATE_BUSY;

        if(cmd==VNA_CMD_SCAN)
        {
            strcpy(vna_status,"Scanning...");
            if(vna_scan_raw())
            {
                strcpy(vna_status,"Scan completed");
            }
        }
        else if(cmd>=VNA_CMD_CAL_OPEN && cmd<=VNA_CMD_CAL_THRU)
        {
            vna_capture(cmd);
        }
        else if(cmd==VNA_CMD_CAL_DONE)
        {
            vna_compute_calibration();
        }
        if(vna_com->open) vna_state=VNA_STATE_CONNECTED;
        else vna_state=VNA_STATE_ERROR;
    }
}
