#!/usr/bin/env python3
"""
JUNO PMT ID bit-encoding utilities.

Each pmt_id is a 32-bit integer encoding:

    id = (1 << 28)
       | (LargeOrSmall << 27)
       | (NorthOrSouth  << 26)
       | (CircleNumber  << 20)
       | (PositionNumber << 12)
       | (UpOrDown      << 11)
       | (PmtType       << 8)

Usage:
    from pmt_info import pmt_type, pmt_id_is_dynode, pmt_id_is_mcp, ...

    pmt_type(363450880)          # --> 2
    pmt_id_is_mcp(363450880)     # --> True
    circle_number(400584960)     # --> ...
"""

# -- Bit positions -----------------------------------------------
_PMT_TYPE_BIT       = 8
_UP_OR_DOWN_BIT     = 11
_POSITION_NUMBER_BIT = 12
_CIRCLE_NUMBER_BIT  = 20
_NORTH_OR_SOUTH_BIT = 26
_LARGE_OR_SMALL_BIT = 27
_MAGIC_BIT          = 28


# -- Raw field extractors ----------------------------------------

def pmt_type(pmt_id):
    """Return PmtType (bits 8-10): 1 = dynode, 2 = MCP, 4 = 3-inch."""
    return (pmt_id >> _PMT_TYPE_BIT) & 0x7


def upp_or_down(pmt_id):
    """Return UpOrDown (bit 11): 0 = up, 1 = down."""
    return (pmt_id >> _UP_OR_DOWN_BIT) & 0x1


def position_number(pmt_id):
    """Return PositionNumber (bits 12-19)."""
    return (pmt_id >> _POSITION_NUMBER_BIT) & 0xFF


def circle_number(pmt_id):
    """Return CircleNumber (bits 20-25)."""
    return (pmt_id >> _CIRCLE_NUMBER_BIT) & 0x3F


def north_or_south(pmt_id):
    """Return NorthOrSouth (bit 26): 0 = south, 1 = north."""
    return (pmt_id >> _NORTH_OR_SOUTH_BIT) & 0x1


def large_or_small(pmt_id):
    """Return LargeOrSmall (bit 27): 0 = 20-inch, 1 = 3-inch."""
    return (pmt_id >> _LARGE_OR_SMALL_BIT) & 0x1


def magic_bit(pmt_id):
    """Return magic bit (bit 28), should always be 1."""
    return (pmt_id >> _MAGIC_BIT) & 0x1


# -- Convenience predicates --------------------------------------

def pmt_is_20inch(pmt_id):
    """True if a 20-inch PMT."""
    return large_or_small(pmt_id) == 0


def pmt_is_3inch(pmt_id):
    """True if a 3-inch PMT."""
    return large_or_small(pmt_id) == 1


def pmt_is_dynode(pmt_id):
    """True if a Hamamatsu R12860 dynode PMT (PmtType = 1)."""
    return pmt_type(pmt_id) == 1


def pmt_is_mcp(pmt_id):
    """True if an NNVT GDB-6201 MCP-PMT (PmtType = 2)."""
    return pmt_type(pmt_id) == 2


def pmt_is_small(pmt_id):
    """True if a 3-inch small PMT (PmtType = 4)."""
    return pmt_type(pmt_id) == 4


def pmt_type_name(pmt_id):
    """Return human-readable PMT type string."""
    mapping = {1: 'dynode (Hamamatsu R12860)', 2: 'MCP (NNVT GDB-6201)', 4: '3-inch small PMT'}
    return mapping.get(pmt_type(pmt_id), 'unknown')


# -- Full decode -------------------------------------------------

def decode(pmt_id):
    """Return a dict of all decoded fields for a pmt_id."""
    return {
        'pmt_id':        pmt_id,
        'pmt_id_hex':    '0x%08X' % pmt_id,
        'pmt_type':      pmt_type(pmt_id),
        'pmt_type_name': pmt_type_name(pmt_id),
        'large_or_small': large_or_small(pmt_id),
        'size':          '20-inch' if pmt_is_20inch(pmt_id) else '3-inch',
        'north_or_south': north_or_south(pmt_id),
        'hemisphere':    'north' if north_or_south(pmt_id) else 'south',
        'circle_number': circle_number(pmt_id),
        'position_number': position_number(pmt_id),
        'upp_or_down':   upp_or_down(pmt_id),
        'magic_bit':     magic_bit(pmt_id),
        'channel':       pmt_id_to_channel(pmt_id),
    }


# -- Channel number <-> pmt_id mapping (via junoid.pq) ------------

_CHANNEL_MAP_FILE = '/junofs/C/junoid/junoid.pq'
_ch_to_pmt_id = None
_pmt_id_to_ch = None


def _load_channel_map():
    """Lazy-load the ch <-> pmt_id mapping from junoid.pq (positive copy_numbers only)."""
    global _ch_to_pmt_id, _pmt_id_to_ch
    if _ch_to_pmt_id is not None:
        return
    import pyarrow.parquet as pq
    t = pq.ParquetFile(_CHANNEL_MAP_FILE).read_row_group(0)
    copies = t.column('copy_number').to_pylist()
    ids = t.column('id').to_pylist()
    _ch_to_pmt_id = {}
    _pmt_id_to_ch = {}
    for ch, pid in zip(copies, ids):
        if ch is not None and pid is not None and ch >= 0:
            _ch_to_pmt_id[ch] = pid
            _pmt_id_to_ch[pid] = ch


def channel_to_pmt_id(ch):
    """Map channel number (copy_number, 0-based) -> hardware pmt_id."""
    _load_channel_map()
    return _ch_to_pmt_id.get(ch)


def pmt_id_to_channel(pmt_id):
    """Map hardware pmt_id -> channel number (copy_number, 0-based)."""
    _load_channel_map()
    return _pmt_id_to_ch.get(pmt_id)


# -- CLI ---------------------------------------------------------
if __name__ == '__main__':
    import sys

    if len(sys.argv) < 2:
        # Demo with a few example PMTs
        examples = [363450880, 269771264, 400584960, 772300800]
        for pid in examples:
            d = decode(pid)
            print('pmt_id=%d (%s)  type=%s  size=%s  hemisphere=%s  circle=%d  pos=%d  updown=%d' % (
                d['pmt_id'], d['pmt_id_hex'], d['pmt_type_name'], d['size'],
                d['hemisphere'], d['circle_number'], d['position_number'], d['upp_or_down']))
    else:
        for arg in sys.argv[1:]:
            pid = int(arg)
            d = decode(pid)
            print()
            for k, v in d.items():
                print('  %-16s %s' % (k, v))