#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy>=1.26", "thrift>=0.20"]
# ///
"""Convert MoSPI Time Use Survey .Nesstar files to compressed, full-size Parquet.

Usage: uv run tus_nesstar_to_parquet.py TUS2024.Nesstar -o output/

Reads column offsets from the NESSTART resource index; does not need CSV,
PyArrow, or a separate metadata JSON. Memory-maps the original input file.
"""

from __future__ import annotations

import argparse
import gzip
import mmap
import os
from pathlib import Path
import struct
import time

import numpy as np
from thrift.protocol import TCompactProtocol
from thrift.transport import TTransport
from thrift.Thrift import TType

NUM = 2
TXT = 6


def read_datasets(path: Path) -> list[dict]:
    """Read data-block descriptors, using *physical* string widths from the index.

    MoSPI TUS declares width 3 for Stratum and Sub_Stratum, but stores 2 bytes
    per row. The indexed payload length, not the declared width, is correct.
    """
    def u16(data, offset):
        return struct.unpack_from('<H', data, offset)[0]

    def u32(data, offset):
        return struct.unpack_from('<I', data, offset)[0]

    with path.open('rb') as f, mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as data:
        if data[:8] != b'NESSTART':
            raise ValueError(f'{path}: not a supported NESSTART container')
        index_offset = u32(data, 0x25)
        count = u32(data, index_offset)
        if index_offset + 4 + 15 * count > len(data):
            raise ValueError('Truncated Nesstar resource index')
        index = {}
        for i in range(count):
            p = index_offset + 4 + 15 * i
            index[u32(data, p)] = (u32(data, p + 4), u32(data, p + 10))

        descriptor_offset = index[u32(data, 0x2F)][0]
        dataset_count = data[0x2B]
        descriptor_width = u16(data, 0x2D)
        datasets = []
        for i in range(dataset_count):
            p = descriptor_offset + descriptor_width * i
            dataset_number = u32(data, p)
            variable_count = u32(data, p + 4)
            rows = u32(data, p + 8)
            if rows != u32(data, p + 12):
                raise ValueError(f'Dataset {dataset_number}: conflicting row counts')
            entry_width = u16(data, p + 20)
            directory_offset = index[u32(data, p + 22)][0]
            variables = []
            for j in range(variable_count):
                q = directory_offset + entry_width * j
                variable_id = u32(data, q + 15)
                name = data[q + 63:q + 127].decode('utf-16le').split('\x00', 1)[0]
                offset, size = index[variable_id]
                mode = data[q + 159]
                fmt = data[q + 5]
                declared_width = data[q + 149]
                if offset + size > len(data) or (mode not in (1, 5)):
                    raise ValueError(f'Unsupported or truncated variable: {name}')
                if mode == 1:
                    if not rows or size % rows:
                        raise ValueError(f'Invalid string payload length: {name}')
                    width = size // rows
                    if width != declared_width:
                        print(f'  {name}: stored width {width}, declared {declared_width}; using stored', flush=True)
                else:
                    width = declared_width
                    expected = {2: (rows + 1) // 2, 3: rows, 4: rows * 2,
                                5: rows * 3, 6: rows * 4, 7: rows * 5,
                                10: rows * 8}.get(fmt)
                    if expected is None or abs(size - expected) > 1:
                        raise ValueError(f'Unsupported numeric encoding: {name} (format={fmt})')
                variables.append({
                    'name': name,
                    'mode': mode,
                    'format': fmt,
                    'width': width,
                    'offset': struct.unpack_from('<q', data, q + 6)[0],
                    'payload_start': offset,
                    'payload_size': size,
                })
            datasets.append({'dataset_number': dataset_number,
                             'record_count': rows, 'variables': variables})
        return datasets


def write_struct(p, fields):
    p.writeStructBegin('')
    for ident, typ, val in fields:
        p.writeFieldBegin('',typ,ident)
        if typ==TType.STRUCT: write_struct(p,val)
        elif typ==TType.LIST:
            elem_typ, seq=val
            p.writeListBegin(elem_typ,len(seq))
            for el in seq:
                if elem_typ==TType.STRUCT: write_struct(p,el)
                elif elem_typ==TType.STRING: p.writeString(el)
                elif elem_typ==TType.I32:p.writeI32(el)
                elif elem_typ==TType.I64:p.writeI64(el)
                else:raise NotImplementedError(elem_typ)
            p.writeListEnd()
        elif typ==TType.STRING: p.writeString(val)
        elif typ==TType.I32:p.writeI32(val)
        elif typ==TType.I64:p.writeI64(val)
        else: raise NotImplementedError(typ)
        p.writeFieldEnd()
    p.writeFieldStop();p.writeStructEnd()

def thrift_bytes(fields):
    buf=TTransport.TMemoryBuffer();p=TCompactProtocol.TCompactProtocol(buf);write_struct(p,fields)
    return buf.getvalue()

def U32(a): return np.asarray(a,dtype='<u4')

def lengths_encode(arr, width):
    """Vectorized plain-encoding of UTF8 strings, dropping NUL terminators."""
    n=len(arr)
    raw=arr.view('u1').reshape(n,width)
    # All missing values are NUL-padded. Numeric-coded strings are ASCII.
    nonzero=(raw!=0)
    lengths=nonzero.sum(axis=1).astype(np.uint32)
    present=lengths>0
    if np.all(lengths[present]==width):
        # Common case: either full-width codes or all-NUL missing values.
        src=arr[present]
        out=np.empty(len(src),dtype=[('size','<u4'),('data',f'S{width}')])
        out['size']=width;out['data']=src
        return present,out.tobytes()
    # Several variables include shorter NUL-terminated code strings.
    src=raw[present];lens=lengths[present]
    start=np.cumsum(lens.astype(np.int64)+4)-lens.astype(np.int64)-4
    out=np.empty(int((lens.astype(np.int64)+4).sum()),dtype='u1')
    for k in range(4):out[start+k]=(lens>>(8*k))&255
    for j in range(width):
        valid=lens>j
        out[start[valid]+4+j]=src[valid,j]
    return present,out.tobytes()

def numeric_decode(d,v,row,n):
    start=v['payload_start'];f=v['format'];bias=v['offset']
    sizes={2:0,3:1,4:2,5:3,6:4,7:5,10:8};w=sizes[f]
    if f==2:
        raw=np.frombuffer(d,dtype='u1',count=(row+n+1)//2,offset=start)
        positions=np.arange(row,row+n)
        x=raw[positions//2].astype(np.uint64)
        x=np.where(positions%2==0, x>>4, x&15)
        sentinel=15
    elif f in (3,4,6):
        dt={3:'u1',4:'<u2',6:'<u4'}[f]
        x=np.frombuffer(d,dtype=dt,count=n,offset=start+row*w).astype(np.uint64)
        sentinel=(1 << (8*w))-1
    elif f in (5,7):
        src=np.frombuffer(d,dtype='u1',count=n*w,offset=start+row*w).reshape(n,w)
        x=np.zeros(n,dtype=np.uint64)
        for k in range(w): x|=src[:,k].astype(np.uint64) << (8*k)
        sentinel=(1 << (8*w))-1
    elif f==10:
        x=np.frombuffer(d,dtype='<f8',count=n,offset=start+row*8).copy()
        mask=~np.isinf(x)
        return mask, x[mask].astype('<f8').tobytes() # float64
    else:raise ValueError((v['name'],f))
    present=x != sentinel
    x=x[present].astype(np.int64)
    if bias:x+=bias
    return present,x.astype('<i8').tobytes()

def page_payload(d,v,row,n):
    if v['mode']==5:return numeric_decode(d,v,row,n)
    w=v['width']; start=v['payload_start']+row*w
    arr=np.frombuffer(d,dtype=f'S{w}',count=n,offset=start)
    return lengths_encode(arr,w)

def vbyte(val):
    b=bytearray()
    while val>=128:b.append((val&127)|128);val>>=7
    b.append(val);return bytes(b)

def definition_levels(present):
    # Encoded as RLE for bit_width=1, with 4-byte length prefix per V1 page.
    if len(present)==0:return b'\0\0\0\0'
    boundaries=np.flatnonzero(present[1:]!=present[:-1])+1
    starts=np.concatenate(([0],boundaries))
    ends=np.concatenate((boundaries,[len(present)]))
    runs=bytearray()
    for s,e in zip(starts,ends):
        runs.extend(vbyte((int(e)-int(s))<<1));runs.append(int(present[s]))
    return struct.pack('<I',len(runs))+runs

def make_page_header(n, ulen, clen):
    dataph=[(1,TType.I32,n),(2,TType.I32,0),(3,TType.I32,3),(4,TType.I32,3)]
    return thrift_bytes([(1,TType.I32,0),(2,TType.I32,ulen),(3,TType.I32,clen),(5,TType.STRUCT,dataph)])

def schema_fields(vars):
    columns=[]
    for v in vars:
        is_number=v['mode']==5
        physical_type=5 if is_number and v['format']==10 else (NUM if is_number else TXT)
        fields=[(1,TType.I32,physical_type),(3,TType.I32,1),(4,TType.STRING,v['name'])]
        if not is_number: fields.append((6,TType.I32,0)) # UTF8 logical converted type
        columns.append(fields)
    return [[(4,TType.STRING,'schema'),(5,TType.I32,len(vars))]]+columns

def field_meta(v,offset,total_compressed,total_uncompressed,n):
    typ=5 if v['mode']==5 and v['format']==10 else NUM if v['mode']==5 else TXT
    return [(1,TType.I32,typ),(2,TType.LIST,(TType.I32,[0,3])),
            (3,TType.LIST,(TType.STRING,[v['name']])),
            (4,TType.I32,2),(5,TType.I64,n),
            (6,TType.I64,total_uncompressed),(7,TType.I64,total_compressed),
            (9,TType.I64,offset)]

def parquet_file(src, dataset, target, row_group_size=500_000, page_size=50_000, max_rows=None):
    ds=dataset; vars=ds['variables']; total=min(ds['record_count'], max_rows or ds['record_count'])
    tic=time.time()
    groups=[]
    with open(src,'rb') as src_file, mmap.mmap(src_file.fileno(),0,access=mmap.ACCESS_READ) as data, open(target,'wb') as f:
        f.write(b'PAR1')
        for group_start in range(0,total,row_group_size):
            m=min(row_group_size,total-group_start)
            columns=[];group_usize=0
            for ci,var in enumerate(vars):
                block_start=f.tell();uncompressed=0
                for page_start in range(group_start,group_start+m,page_size):
                    k=min(page_size,group_start+m-page_start)
                    present,encoded=page_payload(data,var,page_start,k)
                    payload=definition_levels(present)+encoded
                    # Parquet V1 compression encodes definition levels and values together.
                    compressed=gzip.compress(payload,compresslevel=2,mtime=0)
                    header_bytes=make_page_header(k,len(payload),len(compressed))
                    f.write(header_bytes)
                    f.write(compressed)
                    uncompressed+=len(payload)+len(header_bytes)
                block_csize=f.tell()-block_start
                group_usize+=uncompressed
                columns.append([(2,TType.I64,block_start),(3,TType.STRUCT,field_meta(var,block_start,block_csize,uncompressed,m))])
            groups.append([(1,TType.LIST,(TType.STRUCT,columns)),(2,TType.I64,group_usize),(3,TType.I64,m)])
            f.flush()
            print(f'  {target.name} {group_start+m:,}/{total:,} rows ({f.tell()/1048576:.1f} MiB, {time.time()-tic:.1f}s)',flush=True)
        footer=thrift_bytes([(1,TType.I32,1),(2,TType.LIST,(TType.STRUCT,schema_fields(vars))),
                             (3,TType.I64,total),(4,TType.LIST,(TType.STRUCT,groups)),
                             (6,TType.STRING,'TUS2024 Nesstar reader; Python Thrift Parquet writer')])
        f.write(footer);f.write(struct.pack('<I',len(footer)));f.write(b'PAR1')
    return {'rows':total,'columns':len(vars),'bytes':target.stat().st_size,'seconds':round(time.time()-tic,1)}

def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('nesstar', type=Path, help='Input .Nesstar file')
    parser.add_argument('-o', '--output-dir', type=Path, default=Path('.'), help='Output folder')
    parser.add_argument('--max-rows', type=int, help='Convert only this many rows per dataset (for testing)')
    args = parser.parse_args()
    if args.max_rows is not None and args.max_rows < 1:
        parser.error('--max-rows must be positive')
    args.output_dir.mkdir(parents=True, exist_ok=True)

    for dataset in read_datasets(args.nesstar):
        number = dataset['dataset_number']
        # TUS 2024 labels: dataset 2 = activities; dataset 3 = households.
        suffix = {2: 'activities', 3: 'households'}.get(number, f'dataset{number:02}')
        output = args.output_dir / f'{args.nesstar.stem}_{suffix}.parquet'
        partial = output.with_suffix('.parquet.part')
        try:
            print(f'Converting dataset {number} to {output} ...', flush=True)
            stats = parquet_file(args.nesstar, dataset, partial, max_rows=args.max_rows)
            os.replace(partial, output)  # Publish only a complete Parquet file.
        finally:
            partial.unlink(missing_ok=True)
        print(f'Done: {output} — {stats["rows"]:,} rows, '
              f'{stats["columns"]} columns, {stats["bytes"]:,} bytes', flush=True)


if __name__ == '__main__':
    main()