#!/usr/bin/env python3

import argparse
import glob
import json
import mimetypes
import os
import sys
from typing import Dict, Tuple

import tqdm
import skimage.draw
import numpy as np
import imageio
import imageio.plugins.ffmpeg
import cv2

from deface import __version__
from deface.centerface import CenterFace
from deface.sort import Sort

# TODO: Optionally preserve audio track?

def scale_bb(x1, y1, x2, y2, mask_scale=1.0):
    s = mask_scale - 1.0
    h, w = y2 - y1, x2 - x1
    y1 -= h * s
    y2 += h * s
    x1 -= w * s
    x2 += w * s
    return np.round([x1, y1, x2, y2]).astype(int)

def draw_det(
        frame, score, det_idx, x1, y1, x2, y2,
        replacewith: str = 'blur',
        ellipse: bool = True,
        draw_scores: bool = False,
        ovcolor: Tuple[int] = (0, 0, 0)
):
    if replacewith == 'solid':
        cv2.rectangle(frame, (x1, y1), (x2, y2), ovcolor, -1)
    elif replacewith == 'blur':
        bf = 2  # blur factor (number of pixels in each dimension that the face will be reduced to)
        blurred_box =  cv2.blur(
            frame[y1:y2, x1:x2],
            (abs(x2 - x1) // bf, abs(y2 - y1) // bf)
        )
        if ellipse:
            roibox = frame[y1:y2, x1:x2]
            # Get y and x coordinate lists of the "bounding ellipse"
            ey, ex = skimage.draw.ellipse((y2 - y1) // 2, (x2 - x1) // 2, (y2 - y1) // 2, (x2 - x1) // 2)
            roibox[ey, ex] = blurred_box[ey, ex]
            frame[y1:y2, x1:x2] = roibox
        else:
            frame[y1:y2, x1:x2] = blurred_box
    elif replacewith == 'none':
        pass
    if draw_scores:
        cv2.putText(
            frame, f'{score:.2f}', (x1 + 0, y1 - 20),
            cv2.FONT_HERSHEY_DUPLEX, 0.5, (0, 255, 0)
        )

def anonymize_frame(
        dets, frame, mask_scale,
        replacewith, ellipse, draw_scores
):
    for i, det in enumerate(dets):
        boxes, score = det[:4], det[4]
        x1, y1, x2, y2 = boxes.astype(int)
        x1, y1, x2, y2 = scale_bb(x1, y1, x2, y2, mask_scale)
        # Clip bb coordinates to valid frame region
        y1, y2 = max(0, y1), min(frame.shape[0] - 1, y2)
        x1, x2 = max(0, x1), min(frame.shape[1] - 1, x2)

        draw_det(
            frame, score, i, x1, y1, x2, y2,
            replacewith=replacewith,
            ellipse=ellipse,
            draw_scores=draw_scores,
        )

def cam_read_iter(reader):
    while True:
        yield reader.get_next_data()

def video_detect(
        ipath: str,
        opath: str,
        centerface: str,
        threshold: float,
        enable_preview: bool,
        cam: bool,
        nested: bool,
        replacewith: str,
        mask_scale: float,
        ellipse: bool,
        draw_scores: bool,
        ffmpeg_config: Dict[str, str],
        tracking_kwargs: Dict[str, float],
):
    try:
        reader: imageio.plugins.ffmpeg.FfmpegFormat.Reader = imageio.get_reader(ipath)
        meta = reader.get_meta_data()
        _ = meta['size']
    except:
        if cam:
            print(f'Could not find video device {ipath}. Please set a valid input.')
        else:
            print(f'Could not open file {ipath} as a video file with imageio. Skipping file...')
        return

    if cam:
        nframes = None
        read_iter = cam_read_iter(reader)
    else:
        read_iter = reader.iter_data()
        nframes = reader.count_frames()
    if nested:
        bar = tqdm.tqdm(dynamic_ncols=True, total=nframes, position=1, leave=True)
    else:
        bar = tqdm.tqdm(dynamic_ncols=True, total=nframes)

    if opath is not None:
        writer: imageio.plugins.ffmpeg.FfmpegFormat.Writer = imageio.get_writer(
            opath, format='FFMPEG', mode='I', fps=meta['fps'], **ffmpeg_config
        )

    if tracking_kwargs is not None:
        # use SORT BBox tracker from: https://github.com/abewley/sort
        tracker = Sort(**tracking_kwargs)

    for frame in read_iter:
        # Perform network inference, get bb dets but discard landmark predictions
        dets, _ = centerface(frame, threshold=threshold)

        if tracking_kwargs is not None:
            dets = tracker.update(dets)

        anonymize_frame(
            dets, frame, mask_scale=mask_scale,
            replacewith=replacewith, ellipse=ellipse, draw_scores=draw_scores
        )

        if opath is not None:
            writer.append_data(frame)

        if enable_preview:
            cv2.imshow('Preview of anonymization results (quit by pressing Q or Escape)', frame[:, :, ::-1])  # RGB -> RGB
            if cv2.waitKey(1) & 0xFF in [ord('q'), 27]:  # 27 is the escape key code
                cv2.destroyAllWindows()
                break
        bar.update()
    reader.close()
    if opath is not None:
        writer.close()
    bar.close()

def image_detect(
        ipath: str,
        opath: str,
        centerface: str,
        threshold: float,
        replacewith: str,
        mask_scale: float,
        ellipse: bool,
        draw_scores: bool,
        enable_preview: bool
):
    frame = imageio.imread(ipath)

    # Perform network inference, get bb dets but discard landmark predictions
    dets, _ = centerface(frame, threshold=threshold)

    anonymize_frame(
        dets, frame, mask_scale=mask_scale,
        replacewith=replacewith, ellipse=ellipse, draw_scores=draw_scores
    )

    if enable_preview:
        cv2.imshow('Preview of anonymization results (quit by pressing Q or Escape)', frame[:, :, ::-1])  # RGB -> RGB
        if cv2.waitKey(0) & 0xFF in [ord('q'), 27]:  # 27 is the escape key code
            cv2.destroyAllWindows()

    imageio.imsave(opath, frame)
    # print(f'Output saved to {opath}')

def get_file_type(path):
    if path.startswith('<video'):
        return 'cam'
    if not os.path.isfile(path):
        return 'notfound'
    mime = mimetypes.guess_type(path)[0]
    if mime is None:
        return None
    if mime.startswith('video'):
        return 'video'
    if mime.startswith('image'):
        return 'image'
    return mime

def parse_cli_args():
    parser = argparse.ArgumentParser(description='Video anonymization by face detection', add_help=False)
    parser.add_argument(
        'input', nargs='*',
        help=f'File path(s) or camera device name. It is possible to pass multiple paths by separating them by spaces or by using shell expansion (e.g. `$ deface vids/*.mp4`). If a camera is installed, a live webcam demo can be started by running `$ deface cam` (which is a shortcut for `$ deface -p \'<video0>\'`.')
    parser.add_argument(
        '--output', '-o', default=None, metavar='O',
        help='Output file name. Defaults to input path + postfix "_anonymized".')
    parser.add_argument(
        '--thresh', '-t', default=0.2, type=float, metavar='T',
        help='Detection threshold (tune this to trade off between false positive and false negative rate). Default: 0.2.')
    parser.add_argument(
        '--scale', '-s', default=None, metavar='WxH',
        help='Downscale images for network inference to this size (format: WxH, example: --scale 640x360).')
    parser.add_argument(
        '--preview', '-p', default=False, action='store_true',
        help='Enable live preview GUI (can decrease performance).')
    parser.add_argument(
        '--boxes', default=False, action='store_true',
        help='Use boxes instead of ellipse masks.')
    parser.add_argument(
        '--draw-scores', default=False, action='store_true',
        help='Draw detection scores onto outputs.')
    parser.add_argument(
        '--mask-scale', default=1.3, type=float, metavar='M',
        help='Scale factor for face masks, to make sure that masks cover the complete face. Default: 1.3.')
    parser.add_argument(
        '--replacewith', default='blur', choices=['blur', 'solid', 'none'],
        help='Anonymization filter mode for face regions. "blur" applies a strong gaussian blurring, "solid" draws a solid black box and "none" does leaves the input unchanged. Default: "blur".')
    parser.add_argument(
        '--ffmpeg-config', default={"codec": "libx264"}, type=json.loads,
        help='FFMPEG config arguments for encoding output videos. This argument is expected in JSON notation. For a list of possible options, refer to the ffmpeg-imageio docs. Default: \'{"codec": "libx264"}\'.'
    )  # See https://imageio.readthedocs.io/en/stable/format_ffmpeg.html#parameters-for-saving
    parser.add_argument(
        '--track-config', default=None, type=json.loads,
        help='SORT arguments for bounding-box detection tracking, or None if tracking should not be used. This argument is expected in JSON notation. For a list of possible options, refer to the original SORT repository. Default: None.'
    )   # see https://github.com/abewley/sort/blob/master/sort.py#L199
    parser.add_argument(
        '--backend', default='auto', choices=['auto', 'onnxrt', 'opencv'],
        help='Backend for ONNX model execution. Default: "auto" (prefer onnxrt if available).')
    parser.add_argument(
        '--version', action='version', version=__version__,
        help='Print version number and exit.')
    parser.add_argument('--help', '-h', action='help', help='Show this help message and exit.')

    args = parser.parse_args()

    if len(args.input) == 0:
        parser.print_help()
        print('\nPlease supply at least one input path.')
        exit(1)

    if args.input == ['cam']:  # Shortcut for webcam demo with live preview
        args.input = ['<video0>']
        args.preview = True

    return args

def main():
    args = parse_cli_args()
    ipaths = args.input

    base_opath = args.output
    replacewith = args.replacewith
    enable_preview = args.preview
    draw_scores = args.draw_scores
    threshold = args.thresh
    ellipse = not args.boxes
    mask_scale = args.mask_scale
    ffmpeg_config = args.ffmpeg_config
    backend = args.backend
    in_shape = args.scale
    track_config = args.track_config
    if in_shape is not None:
        w, h = in_shape.split('x')
        in_shape = int(w), int(h)

    # TODO: scalar downscaling setting (-> in_shape), preserving aspect ratio
    centerface = CenterFace(in_shape=in_shape, backend=backend)

    multi_file = len(ipaths) > 1
    if multi_file:
        ipaths = tqdm.tqdm(ipaths, position=0, dynamic_ncols=True, desc='Batch progress')

    for ipath in ipaths:
        opath = base_opath
        if ipath == 'cam':
            ipath = '<video0>'
            enable_preview = True
        filetype = get_file_type(ipath)
        is_cam = filetype == 'cam'
        if opath is None and not is_cam:
            root, ext = os.path.splitext(ipath)
            opath = f'{root}_anonymized{ext}'
        print(f'Input:  {ipath}\nOutput: {opath}')
        if opath is None and not enable_preview:
            print('No output file is specified and the preview GUI is disabled. No output will be produced.')
        if filetype == 'video' or is_cam:
            video_detect(
                ipath=ipath,
                opath=opath,
                centerface=centerface,
                threshold=threshold,
                cam=is_cam,
                replacewith=replacewith,
                mask_scale=mask_scale,
                ellipse=ellipse,
                draw_scores=draw_scores,
                enable_preview=enable_preview,
                nested=multi_file,
                ffmpeg_config=ffmpeg_config,
                tracking_kwargs=track_config
            )
        elif filetype == 'image':
            image_detect(
                ipath=ipath,
                opath=opath,
                centerface=centerface,
                threshold=threshold,
                replacewith=replacewith,
                mask_scale=mask_scale,
                ellipse=ellipse,
                draw_scores=draw_scores,
                enable_preview=enable_preview
            )
        elif filetype is None:
            print(f'Can\'t determine file type of file {ipath}. Skipping...')
        elif filetype == 'notfound':
            print(f'File {ipath} not found. Skipping...')
        else:
            print(f'File {ipath} has an unknown type {filetype}. Skipping...')

if __name__ == '__main__':
    main()


# --- Do not remove these libs ---
from freqtrade.strategy.interface import IStrategy
from typing import Dict, List
from functools import reduce
from pandas import DataFrame, Series
# --------------------------------

import logging
import pandas as pd
import numpy as np
import datetime
from freqtrade.persistence import Trade

logger = logging.getLogger(__name__)

class YourStrat(IStrategy):
    # replace this by your strategy
    pass

class TrailingBuyStrat(YourStrat):
    # This class is designed to heritate from yours and starts trailing buy with your buy signals
    # Trailing buy starts at any buy signal
    # Trailing buy stops  with BUY if : price decreases and rises again more than trailing_buy_offset
    # Trailing buy stops with NO BUY : current price is > intial price * (1 +  trailing_buy_max) OR custom_sell tag
    # IT IS NOT COMPATIBLE WITH BACKTEST/HYPEROPT
    #
    # if process_only_new_candles = True, then you need to use 1m timeframe (and normal strat timeframe as informative)
    # if process_only_new_candles = False, it will use ticker data and you won't need to change anything

    trailing_buy_order_enabled = True
    trailing_buy_offset = 0.005  # rebound limit before a buy in % of initial price
    # (example with 0.5%. initial price : 100 (uplimit is 100.5), 2nd price : 99 (no buy, uplimit updated to 99.5), 3price 98 (no buy uplimit updated to 98.5), 4th price 99 -> BUY
    trailing_buy_max = 0.1  # stop trailing buy if current_price > starting_price * (1+trailing_buy_max)

    process_only_new_candles = False

    custom_info = dict()

    init_trailing_dict = {
        'trailing_buy_order_started': False,
        'trailing_buy_order_uplimit': 0,
        'start_trailing_price': 0,
        'buy_tag': None
    }

    def custom_sell(self, pair: str, trade: Trade, current_time: datetime, current_rate: float,
                    current_profit: float, **kwargs):
        tag = super().custom_sell(pair, trade, current_time, current_rate, current_profit, **kwargs)
        if tag:
            self.custom_info[pair]['trailing_buy'] = self.init_trailing_dict
            logger.info(f'STOP trailing buy for {pair} because of {tag}')
        return tag

    def populate_indicators(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        dataframe = super().populate_indicators(dataframe, metadata)
        if not metadata["pair"] in self.custom_info:
            self.custom_info[metadata["pair"]] = dict()
        if not 'trailing_buy' in self.custom_info[metadata['pair']]:
            self.custom_info[metadata["pair"]]['trailing_buy'] = self.init_trailing_dict
        return dataframe

    def confirm_trade_exit(self, pair: str, trade: Trade, order_type: str, amount: float,
                           rate: float, time_in_force: str, sell_reason: str, **kwargs) -> bool:
        val = super().confirm_trade_exit(pair, trade, order_type, amount, rate, time_in_force, sell_reason, **kwargs)
        self.custom_info[pair]['trailing_buy'] = self.init_trailing_dict
        return val

    def populate_buy_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        def get_local_min(x):
            win = dataframe.loc[:, 'barssince_last_buy'].iloc[x.shape[0] - 1].astype('int')
            win = max(win, 0)
            return pd.Series(x).rolling(window=win).min().iloc[-1]

        dataframe = super().populate_buy_trend(dataframe, metadata)
        dataframe = dataframe.rename(columns={"buy": "pre_buy"})

        if self.trailing_buy_order_enabled and self.config['runmode'].value in ('live', 'dry_run'):  # trailing live dry ticker, 1m
            last_candle = dataframe.iloc[-1].squeeze()
            if not self.process_only_new_candles:
                current_price = self.get_current_price(metadata["pair"])
            else:
                current_price = last_candle['close']
            dataframe['buy'] = 0
            if not self.custom_info[metadata["pair"]]['trailing_buy']['trailing_buy_order_started'] and last_candle['pre_buy'] == 1:
                self.custom_info[metadata["pair"]]['trailing_buy'] = {
                    'trailing_buy_order_started': True,
                    'trailing_buy_order_uplimit': last_candle['close'],
                    'start_trailing_price': last_candle['close'],
                    'buy_tag': last_candle['buy_tag'] if 'buy_tag' in last_candle else 'buy signal'
                }
                logger.info(f'start trailing buy for {metadata["pair"]} at {last_candle["close"]}')
            elif self.custom_info[metadata["pair"]]['trailing_buy']['trailing_buy_order_started']:
                if current_price < self.custom_info[metadata["pair"]]['trailing_buy']['trailing_buy_order_uplimit']:
                    # update uplimit
                    self.custom_info[metadata["pair"]]['trailing_buy']['trailing_buy_order_uplimit'] = min(current_price * (1 + self.trailing_buy_offset), self.custom_info[metadata["pair"]]['trailing_buy']['trailing_buy_order_uplimit'])
                    logger.info(f'update trailing buy for {metadata["pair"]} at {self.custom_info[metadata["pair"]]["trailing_buy"]["trailing_buy_order_uplimit"]}')
                elif current_price < self.custom_info[metadata["pair"]]['trailing_buy']['start_trailing_price']:
                    # buy ! current price > uplimit but lower thant starting price
                    dataframe.iloc[-1, dataframe.columns.get_loc('buy')] = 1
                    ratio = "%.2f" % ((1 - current_price / self.custom_info[metadata['pair']]['trailing_buy']['start_trailing_price']) * 100)
                    if 'buy_tag' in dataframe.columns:
                        dataframe.iloc[-1, dataframe.columns.get_loc('buy_tag')] = f"{self.custom_info[metadata['pair']]['trailing_buy']['buy_tag']} ({ratio} %)"
                    # stop trailing when buy signal ! prevent from buying much higher price when slot is free
                    self.custom_info[metadata["pair"]]['trailing_buy'] = self.init_trailing_dict
                    logger.info(f'STOP trailing buy for {metadata["pair"]} because I buy it {ratio}')
                elif current_price > (self.custom_info[metadata["pair"]]['trailing_buy']['start_trailing_price'] * (1 + self.trailing_buy_max)):
                    self.custom_info[metadata["pair"]]['trailing_buy'] = self.init_trailing_dict
                    logger.info(f'STOP trailing buy for {metadata["pair"]} because of the price is higher than starting prix * {1 + self.trailing_buy_max}')
                else:
                    logger.info(f'price to high for {metadata["pair"]} at {current_price} vs {self.custom_info[metadata["pair"]]["trailing_buy"]["trailing_buy_order_uplimit"]}')
        elif self.trailing_buy_order_enabled:
            # FOR BACKTEST
            # NOT WORKING
            dataframe.loc[
                (dataframe['pre_buy'] == 1) &
                (dataframe['pre_buy'].shift() == 0)
                , 'pre_buy_switch'] = 1
            dataframe['pre_buy_switch'] = dataframe['pre_buy_switch'].fillna(0)

            dataframe['barssince_last_buy'] = dataframe['pre_buy_switch'].groupby(dataframe['pre_buy_switch'].cumsum()).cumcount()

            # Create integer positions of each row
            idx_positions = np.arange(len(dataframe))
            # "shift" those integer positions by the amount in shift col
            shifted_idx_positions = idx_positions - dataframe["barssince_last_buy"]
            # get the label based index from our DatetimeIndex
            shifted_loc_index = dataframe.index[shifted_idx_positions]
            # Retrieve the "shifted" values and assign them as a new column
            dataframe["close_5m_last_buy"] = dataframe.loc[shifted_loc_index, "close_5m"].values

            dataframe.loc[:, 'close_lower'] = dataframe.loc[:, 'close'].expanding().apply(get_local_min)
            dataframe['close_lower'] = np.where(dataframe['close_lower'].isna() == True, dataframe['close'], dataframe['close_lower'])
            dataframe['close_lower_offset'] = dataframe['close_lower'] * (1 + self.trailing_buy_offset)
            dataframe['trailing_buy_order_uplimit'] = np.where(dataframe['barssince_last_buy'] < 20, pd.DataFrame([dataframe['close_5m_last_buy'], dataframe['close_lower_offset']]).min(), np.nan)

            dataframe.loc[
                (dataframe['barssince_last_buy'] < 20) &  # must buy within last 20 candles after signal
                (dataframe['close'] > dataframe['trailing_buy_order_uplimit'])
                , 'trailing_buy'] = 1

            dataframe['trailing_buy_count'] = dataframe['trailing_buy'].rolling(20).sum()

            dataframe.log[
                (dataframe['trailing_buy'] == 1) &
                (dataframe['trailing_buy_count'] == 1)
                , 'buy'] = 1
        else:  # No buy trailing
            dataframe.loc[
                (dataframe['pre_buy'] == 1)
                , 'buy'] = 1
        return dataframe

    def get_current_price(self, pair: str) -> float:
        ticker = self.dp.ticker(pair)
        current_price = ticker['last']
        return current_price


import numpy as np
import struct
from wand.image import Image
from .base import TextureFormat

def _makeDdsHeader(width, height, numBytes):
    data = (
        0x20534444, # header "DDS "
        124, #DWORD           dwSize;
        0x081007, #DWORD           dwFlags;
        height, #DWORD           dwHeight;
        width, #DWORD           dwWidth;
        numBytes, #DWORD           dwPitchOrLinearSize;
        0, #DWORD           dwDepth;
        1, #DWORD           dwMipMapCount;
        0,0,0,0,0,0,0,0,0,0,0, #DWORD           dwReserved1[11];
        #struct DDS_PIXELFORMAT {
        32,    #DWORD dwSize;
        0x04,    #DWORD dwFlags;
        0x31545844,    #DWORD dwFourCC; "DXT1"
        0,    #DWORD dwRGBBitCount;
        0,    #DWORD dwRBitMask;
        0,    #DWORD dwGBitMask;
        0,    #DWORD dwBBitMask;
        0x00000000,    #DWORD dwABitMask;
        #} ddspf;
        0x1000, #DWORD           dwCaps;
        0, #DWORD           dwCaps2;
        0, #DWORD           dwCaps3;
        0, #DWORD           dwCaps4;
        0, #DWORD           dwReserved2;
    )
    return b''.join(map(lambda d: struct.pack('<I', d), data))

class CMPR(TextureFormat):
    """Compressed CMPR/DDS/BC1/DXT1 format."""
    fmtId        = 0x0E
    name         = "CMPR"
    hasAlpha     = True
    bitsPerPixel = 4
    blockW       = 8
    blockH       = 8

    def __init__(self, width:int, height:int, pixels:bytes):
        # byteswap pixels
        data = np.frombuffer(pixels,dtype=np.uint16)
        data = ((data & 0x00FF) << 8) | ((data & 0xFF00) >> 8)

        #data = np.frombuffer(data,dtype=np.uint8)
        #data = data & 0xFF # stupid hack to make data mutable
        #for i in range(0, len(data), 8):
            #data[i+0], data[i+1] = data[i+1], data[i+0]
            #data[i+2], data[i+3] = data[i+3], data[i+2]
            #data[i+4], data[i+5] = data[i+5], data[i+4]
            #data[i+6], data[i+7] = data[i+7], data[i+6]
            #data[i+4], data[i+5], data[i+6], data[i+7] = \
            #data[i+7], data[i+6], data[i+5], data[i+4]

        # generate DDS header
        data = data.tobytes()
        data = _makeDdsHeader(width, height, len(data))+data
        with open('butt.dds', 'wb') as file: file.write(data)
        img  = Image(blob=data, format='dds')
        data = bytes(img.export_pixels())
        data = np.frombuffer(data,dtype=np.uint32)
        data = np.reshape(data, (height, width))

        #data = data & 0xFFFFFFFF
        #for y in range(height):
        #    for x in range(width):
        #        data[y,x] = 0xFF000000 | ((y&0xF)<<4) | ((x&0xF)<<12)

        super().__init__(width, height, data)
        self.descramble()

    def descramble(self):
        # descramble blocks
        # 0123 4567 89AB CDEF GHIJ KLMN OPQR STUV =>
        # 0123 GHIJ
        # 4567 KLMN
        # 89AB OPQR
        # CDEF STUV
        # x,y = x+(y*4),y/4
        pixelsOut = np.zeros_like(self.pixels)
        idx = 0
        bw, bh = 4, 4
        for y in range(0, self.height, bh):
            for x in range(0, self.width, bw):
                for by in range(bh):
                    for bx in range(bw):
                        sx = (idx %  self.width)
                        sy = (idx // self.width)
                        dx = x+bx
                        dy = y+by
                        #print(x, y, sx, sy, dx, dy)
                        #pixelsOut[dy,dx] = self.pixels[sy,sx]
                        try: pixelsOut[sy,sx] = self.pixels[dy,dx]
                        #pixelsOut[y+by,x+bx] = self.pixels[py,px]
                        except IndexError: pass
                        idx += 1
        self.pixels = pixelsOut

    def toTexture(self) -> bytes:
        """Return buffer of this texture in SFA format."""
        header = self._makeSfaHeader()

        # read pixels and convert to RGB565
        # XXX this can probably be done faster
        buf = np.array(dtype=np.uint16, ndmin=3)
        for y in range(self.height):
            for x in range(self.width):
                rgba = self.pixels[x,y]
                r =  rgba >> 24 # XXX verify order
                g = (rgba >> 16) & 0xFF
                b = (rgba >>  8) & 0xFF
                buf[x,y] = (b >> 3) | ((g >> 2) << 5) | ((r >> 3) << 11)

        # byteswap
        buf = ((buf & 0x00FF) << 8) | ((buf & 0xFF00) >> 8)
        return header + buf.tobytes()

    def unswizzle(self):
        """Un-swizzle this texture from the GameCube pixel order.

        This changes the texture; it does not return a copy.
        """
        pixelsOut = np.zeros_like(self.pixels)
        data = self.pixels.ravel() # to 1D array (should be unravel?)
        #data = self.pixels
        idx = 0
        for y in range(0, self.height, self.blockH):
            for x in range(0, self.width, self.blockW):
                for by in range(0, self.blockH):
                    for bx in range(self.blockW):
                        # notice y,x indexing here.
                        try: pixelsOut[y+by, x+bx] = data[idx]
                        except IndexError: continue
                        idx += 1
        self.pixels = pixelsOut
        #self.descramble()
        return self


import numpy as np
import torch

import utils

import models

class CorefSystem:

    def __init__(self, device, config):
        """
        Parameters
        ----------
        device: str
        config: utils.Config
        """
        self.device = device
        self.config = config

        # Initialize model
        if self.config["model_name"] == "joshi2020":
            self.model = models.Joshi2020(device=device, config=config)
        else:
            raise Exception("Invalid model_name %s" % self.config["model_name"])

        # Show parameter shapes
        utils.writelog("Model parameters:")
        for name, param in self.model.named_parameters():
            utils.writelog("%s: %s" % (name, tuple(param.shape)))

    def load_model(self, path):
        """
        Parameters
        ----------
        path: str
        """
        self.model.load_state_dict(torch.load(path, map_location=torch.device("cpu")), strict=False)

    def save_model(self, path):
        """
        Parameters
        ----------
        path: str
        """
        torch.save(self.model.state_dict(), path)

    def to_gpu(self, device):
        """
        Parameters
        ----------
        device: str
        """
        self.model.to(device)

    def compute_loss(self, data):
        """
        Parameters
        ----------
        data: utils.DataInstance

        Returns
        -------
        torch.Tensor
        """
        # Switch to training mode
        self.model.train()

        # Forward
        _, loss = self.model.forward(data=data, get_loss=True)

        return loss

    def predict(self, data, evaluator=None, gold_clusters=None):
        """
        Parameters
        ----------
        data: utils.DataInstance
        evaluator: CorefEvaluator, default None
        gold_clusters: list[list[(int, int)]], default None

        Returns
        -------
        list[list[(int, int)]]
        CorefEvaluator or None
        """
        if evaluator is None or gold_clusters is None:
            assert evaluator is None and gold_clusters is None

        # Switch to inference mode
        self.model.eval()

        # Forward
        (span_starts, span_ends, antecedent_indices, antecedent_scores), _ \
                = self.model.forward(data=data, get_loss=False)

        span_starts = span_starts.tolist()
        span_ends = span_ends.tolist()
        antecedent_indices = antecedent_indices.tolist()
        antecedent_scores = antecedent_scores.tolist()

        # Get predicted antecedents
        predicted_antecedents = self.get_predicted_antecedents(antecedent_indices=antecedent_indices, antecedent_scores=antecedent_scores)

        # Get clusters
        predicted_clusters, mention_to_predicted = self.get_predicted_clusters(
                                                            span_starts=span_starts,
                                                            span_ends=span_ends,
                                                            predicted_antecedents=predicted_antecedents)

        if evaluator is None:
            return predicted_clusters, None

        # Update evaluator
        # mention_to_predicted = {m: predicted_clusters[cluster_idx] for m, cluster_idx in mention_to_cluster_id.items()}
        gold_clusters = [tuple(tuple(m) for m in cluster) for cluster in gold_clusters]
        mention_to_gold = {m: cluster for cluster in gold_clusters for m in cluster}
        evaluator.update(predicted_clusters, gold_clusters, mention_to_predicted, mention_to_gold)

        return predicted_clusters, evaluator

    def get_predicted_antecedents(self, antecedent_indices, antecedent_scores):
        """
        Parameters
        ----------
        antecedent_indices: list[list[int]]
            shape (n_top_spans, n_ant_spans)
        antecedent_scores: list[list[float]]
            shape (n_top_spans, 1 + n_ant_spans)

        Returns
        -------
        list[int]
        """
        predicted_antecedents = []
        for i, idx in enumerate(np.argmax(antecedent_scores, axis=1) - 1):
            if idx < 0:
                # The dummy antecedent is selected.
                # Since the coreference score to the dummy antecedent is always set to zero,
                # the coreference scores to the non-dummy candidates are all negative.
                predicted_antecedents.append(-1)
            else:
                # The maximum antecedent score is positive,
                # and the selected antecedent is not dummy.
                predicted_antecedents.append(antecedent_indices[i][idx])
        return predicted_antecedents

    def get_predicted_clusters(self, span_starts, span_ends, predicted_antecedents):
        """
        Parameters
        ----------
        span_starts: list[int]
        span_ends: list[int]
        predicted_antecedents: list[int]

        Returns
        -------
        list[list[(int, int)]]
        dict[(int, int), int]
        """
        # Get predicted clusters
        predicted_clusters = [] # list[list[(int, int)]]
        mention_to_cluster_id = {} # dict[(int, int), int]
        for mention_i, antecedent_i in enumerate(predicted_antecedents):
            # No coreference
            if antecedent_i < 0:
                continue

            # Check whether the coreference is valid
            assert antecedent_i < mention_i, f'antecedent (index {antecedent_i}) must appear earlier than span (index {mention_i})'

            # Add antecedent to cluster (if the antecedent is chosen for the first time)
            antecedent = (int(span_starts[antecedent_i]), int(span_ends[antecedent_i])) # Antecedent span
            antecedent_cluster_id = mention_to_cluster_id.get(antecedent, -1) # Cluster ID
            if antecedent_cluster_id == -1:
                # Add a new cluster
                antecedent_cluster_id = len(predicted_clusters) # New cluster ID
                predicted_clusters.append([antecedent]) # Add antecedent to cluster
                mention_to_cluster_id[antecedent] = antecedent_cluster_id
            else:
                # This (antecedent) span is already selected as an antecedent of the previous mention(s)
                pass

            # Add mention to cluster
            mention = (int(span_starts[mention_i]), int(span_ends[mention_i])) # Mention span
            assert not mention in predicted_clusters[antecedent_cluster_id]
            assert not mention in mention_to_cluster_id
            predicted_clusters[antecedent_cluster_id].append(mention) # Add mention to cluster
            mention_to_cluster_id[mention] = antecedent_cluster_id

        predicted_clusters = [tuple(c) for c in predicted_clusters]
        mention_to_predicted = {m: predicted_clusters[cluster_idx] for m, cluster_idx in mention_to_cluster_id.items()}
        return predicted_clusters, mention_to_predicted


#  Version: 2020.02.21
#
#  MIT License
#
#  Copyright (c) 2018 Jiankang Deng and Jia Guo
#
#  Permission is hereby granted, free of charge, to any person obtaining a copy
#  of this software and associated documentation files (the "Software"), to deal
#  in the Software without restriction, including without limitation the rights
#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#  copies of the Software, and to permit persons to whom the Software is
#  furnished to do so, subject to the following conditions:
#
#  The above copyright notice and this permission notice shall be included in all
#  copies or substantial portions of the Software.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#  SOFTWARE.
#

from __future__ import division
import mxnet as mx
import numpy as np
import cv2

__all__ = ['FaceRecognition',
           'arcface_r100_v1', 'arcface_outofreach_v1', 'arcface_mfn_v1',
           'get_arcface']

class FaceRecognition:
    def __init__(self, name, download, param_file):
        self.name = name
        self.download = download
        self.param_file = param_file
        self.image_size = (112, 112)
        if download:
            assert param_file

    def prepare(self, ctx_id):
        if self.param_file:
            pos = self.param_file.rfind('-')
            prefix = self.param_file[0:pos]
            pos2 = self.param_file.rfind('.')
            epoch = int(self.param_file[pos+1:pos2])
            sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
            all_layers = sym.get_internals()
            sym = all_layers['fc1_output']
            if ctx_id>=0:
                ctx = mx.gpu(ctx_id)
            else:
                ctx = mx.cpu()
            model = mx.mod.Module(symbol=sym, context=ctx, label_names = None)
            data_shape = (1,3)+self.image_size
            model.bind(data_shapes=[('data', data_shape)])
            model.set_params(arg_params, aux_params)
            #warmup
            data = mx.nd.zeros(shape=data_shape)
            db = mx.io.DataBatch(data=(data,))
            model.forward(db, is_train=False)
            embedding = model.get_outputs()[0].asnumpy()
            self.model = model
        else:
            pass

    def get_embedding(self, img):
        assert self.param_file and self.model
        assert img.shape[2]==3 and img.shape[0:2]==self.image_size
        data = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        data = np.transpose(data, (2,0,1))
        data = np.expand_dims(data, axis=0)
        data = mx.nd.array(data)
        db = mx.io.DataBatch(data=(data,))
        self.model.forward(db, is_train=False)
        embedding = self.model.get_outputs()[0].asnumpy()
        return embedding

    def compute_sim(self, img1, img2):
        emb1 = self.get_embedding(img1).flatten()
        emb2 = self.get_embedding(img2).flatten()
        from numpy.linalg import norm
        sim = np.dot(emb1, emb2)/(norm(emb1)*norm(emb2))
        return sim

def get_arcface(name, download=True,
               root='~/.insightface/models', **kwargs):
    if not download:
        return FaceRecognition(name, False, None)
    else:
        from .model_store import get_model_file
        _file = get_model_file("arcface_%s"%name, root=root)
        return FaceRecognition(name, True, _file)

def arcface_r100_v1(**kwargs):
    return get_arcface("r100_v1", download=True, **kwargs)

def arcface_mfn_v1(**kwargs):
    return get_arcface("mfn_v1", download=True, **kwargs)

def arcface_outofreach_v1(**kwargs):
    return get_arcface("outofreach_v1", download=False, **kwargs)


# -----------------------------------------------------
# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
# Written by Chao Xu (xuchao.19962007@sjtu.edu.cn)
# -----------------------------------------------------

"""API of yolo detector"""
import os
import sys
sys.path.insert(0, os.path.dirname(__file__))
import platform

import torch
import numpy as np

from yolo_v4.detect import Detector
from yolo_v4.preprocess import prep_image, prep_frame

from detector.apis import BaseDetector

#only windows visual studio 2013 ~2017 support compile c/cuda extensions
#If you force to compile extension on Windows and ensure appropriate visual studio
#is intalled, you can try to use these ext_modules.
if platform.system() != 'Windows':
    from detector.nms import nms_wrapper

class YOLOV4Detector(BaseDetector):
    def __init__(self, cfg, opt=None):
        super(YOLOV4Detector, self).__init__()

        self.detector_cfg = cfg
        self.detector_opt = opt
        self.model_cfg = cfg.get('CONFIG', 'detector/yolo_v4/cfg/yolov4.cfg')
        self.model_weights = cfg.get('WEIGHTS', 'detector/yolo_v4/data/yolov4.weights')

        os.environ['YOLOV4_HOME'] = os.path.dirname(self.model_weights)

        self.inp_dim = cfg.get('INP_DIM', 608)
        self.nms_thresh = cfg.get('NMS_THRES', 0.6)
        self.confidence = cfg.get('CONFIDENCE', 0.05)
        # self.num_classes = cfg.get('NUM_CLASSES', 80)
        self.model = None

    def load_model(self):
        args = self.detector_opt

        print('Loading YOLOv4 model..')
        with torch.cuda.device(args.gpus[0] if len(args.gpus) > 0 else -1):
            self.model = Detector(configfile=self.model_cfg, weightsfile=self.model_weights, conf_threshold=self.confidence, nms_threshold=self.nms_thresh, device_ids=args.gpus, default_device=args.device)
        print("Network successfully loaded")

    def image_preprocess(self, img_source):
        """
        Pre-process the img before fed to the object detection network
        Input: image name(str) or raw image data(ndarray or torch.Tensor,channel GBR)
        Output: pre-processed image data(torch.FloatTensor,(1,3,h,w))
        """
        if isinstance(img_source, str):
            img, orig_img, im_dim_list = prep_image(img_source, self.inp_dim)
        elif isinstance(img_source, torch.Tensor) or isinstance(img_source, np.ndarray):
            img, orig_img, im_dim_list = prep_frame(img_source, self.inp_dim)
        else:
            raise IOError('Unknown image source type: {}'.format(type(img_source)))

        return img

    def images_detection(self, imgs, orig_dim_list):
        """
        Feed the img data into object detection network and
        collect bbox w.r.t original image size
        Input: imgs(torch.FloatTensor,(b,3,h,w)): pre-processed mini-batch image input
               orig_dim_list(torch.FloatTensor, (b,(w,h,w,h))): original mini-batch image size
        Output: dets(torch.cuda.FloatTensor,(n,(batch_idx,x1,y1,x2,y2,c,s,idx of cls))): human detection results
        """
        args = self.detector_opt
        if not self.model:
            self.load_model()

        with torch.no_grad():
            imgs = imgs.to(args.device) if args else imgs.cuda()
            imgs, class_ids, scores, boxes = self.model.detect(imgs)

            dets = None
            write = False
            for idx, _ in enumerate(imgs):
                if len(boxes[idx]) == 0:
                    continue

                det_new = torch.empty([len(boxes[idx]), 8])
                det_new[:, 0] = idx  #index of img
                det_new[:, 1:3] = torch.from_numpy(boxes[idx][:, 0:2])  # bbox x1,y1
                det_new[:, 3:5] = torch.from_numpy(boxes[idx][:, 2:4])  # bbox x2,y2
                det_new[:, 6:7] = torch.from_numpy(scores[idx])  # cls conf
                det_new[:, 7:8] = torch.from_numpy(class_ids[idx])   # cls idx

                if not write:
                    dets = det_new
                    write = True
                else:
                    dets = torch.cat((dets, det_new))

            if dets is None:
                return 0

            dets = dets.cpu()

            w_scaling_factor = (self.inp_dim / orig_dim_list[:, 0]).view(-1, 1)
            h_scaling_factor = (self.inp_dim / orig_dim_list[:, 1]).view(-1, 1)

            w_scale = torch.empty([dets.shape[0], 1], dtype=torch.float32)
            h_scale = torch.empty([dets.shape[0], 1])

            start = 0
            for idx, _ in enumerate(boxes):
                if len(boxes[idx]) > 0:
                    w_scale[start:start+len(boxes[idx])] = w_scaling_factor[idx]
                    h_scale[start:start+len(boxes[idx])] = h_scaling_factor[idx]
                    start += len(boxes[idx])

            dets[:, [1, 3]] /= w_scale
            dets[:, [2, 4]] /= h_scale

            orig_dim_list = torch.index_select(orig_dim_list, 0, dets[:, 0].long())
            for i in range(dets.shape[0]):
                dets[i, [1, 3]] = torch.clamp(dets[i, [1, 3]], 0.0, orig_dim_list[i, 0])
                dets[i, [2, 4]] = torch.clamp(dets[i, [2, 4]], 0.0, orig_dim_list[i, 1])

            return dets

    def detect_one_img(self, img_name):
        """
        Detect bboxs in one image
        Input: 'str', full path of image
        Output: '[{"category_id":1,"score":float,"bbox":[x,y,w,h],"image_id":str},...]',
        The output results are similar with coco results type, except that image_id uses full path str
        instead of coco %012d id for generalization.
        """
        if not self.model:
            self.load_model()

        #pre-process(scale, normalize, ...) the image
        img, orig_img, orig_dim_list = prep_image(img_name, self.inp_dim)

        if isinstance(img, np.ndarray):
            img = torch.from_numpy(img)

        # add one dimension at the front for batch if image shape (3,h,w)
        if img.dim() == 3:
            img = img.unsqueeze(0)

        img_dim_list = torch.FloatTensor([orig_dim_list]).repeat(1, 2)
        dets = self.images_detection(img, img_dim_list).tolist()

        dets_results = []
        for det in dets:
            det_dict = dict()
            det_dict["category_id"] = det[6]
            det_dict["score"] = float(det[5])
            det_dict["bbox"] = [det[1], det[2], det[3]-det[1], det[4]-det[2]]
            det_dict["image_id"] = os.path.basename(img_name)
            dets_results.append(det_dict)

        return dets_results


#!/usr/bin/env python

"""mono.py: PyDDS Monophonic ADSR Synthesis Module
"""

import time as tm
import numpy as np
import matplotlib.pyplot as plt
import pyaudio

__version__     = "0.2.1"
__status__      = "Development"

__author__      = "Randy Rubin"
__copyright__   = "Copyright 2018, Randy Rubin"
__license__     = "MIT"

class DDS:

    _TAB_INDEX_BITS = 16
    _PHASE_ACC_BITS = 32
    _SAMPLE_BITS = 16
    _ENV_VAL_BITS = 16

    _TAB_LEN = int(round(2**_TAB_INDEX_BITS))
    _PHASE_ACC_MAX = int(round(2**_PHASE_ACC_BITS)) - 1
    _PHASE_ACC_SHIFT = _PHASE_ACC_BITS - _TAB_INDEX_BITS
    _SAMPLE_ABSMAX = int(round(2**(_SAMPLE_BITS - 1))) - 1
    _ENV_VAL_ABSMAX = int(round(2**_ENV_VAL_BITS)) - 1

    def __init__(self, VERBOSE=False,
                OUT_BUF_UPDATE_MS=5,
                amp_ratio=0.25,
                freq_hz=440,
                phase_ratio=0,
                attack_ms=50,
                decay_ms=100,
                sustain_ratio=0.5,
                release_ms=333,
                cancel_ms=3,
                retrigger=False,
                recycle=True,
                decimation=False,
                bit_depth=False):

        self._start_time = tm.time()
        self.VERBOSE = VERBOSE

        self._verbose("Initializing synthesis module...")

        self.OUT_BUF_UPDATE_MS = OUT_BUF_UPDATE_MS
        self._init_pyaudio()

        self.set_amp_ratio(amp_ratio)
        self.set_freq_hz(freq_hz)
        self.set_phase_ratio(phase_ratio)
        self.set_attack_ms(attack_ms)
        self.set_decay_ms(decay_ms)
        self.set_sustain_ratio(sustain_ratio)
        self.set_release_ms(release_ms)
        self.set_cancel_ms(cancel_ms)
        self.ena_retrigger() if retrigger else self.dis_retrigger()
        self.ena_recycle() if recycle else self.dis_recycle()
        self.set_decimation(decimation) if decimation else self.dis_decimation()
        self.set_bit_depth(bit_depth) if bit_depth else self.dis_bit_depth()

        self._load_all_tabs()
        self.sel_sin()
        self.note_off()

        self._verbose("Synthesis module initialization complete.")

    def __del__(self):
        self.close()

    def _timestamp(self):
        return tm.time()-self._start_time

    def _verbose(self, msg):
        if self.VERBOSE:
            print("{:.3f} {}".format(self._timestamp(), msg))

    def _init_pyaudio(self):
        self._verbose("Setup audio device interface...")
        self._pa = pyaudio.PyAudio()
        self._DEV_INFO = self._pa.get_default_output_device_info()
        self._SAMPLE_RATE_HZ = int(round(self._DEV_INFO['defaultSampleRate']))
        out_buf_len = int(round(self._SAMPLE_RATE_HZ * (self.OUT_BUF_UPDATE_MS / 1000)))
        self._out_buf = np.zeros(out_buf_len, dtype=np.int16)
        self._verbose("Audio device interface is ready.")

    def _pyaudio_callback(self, in_data, frame_count, time_info, status):
        self._load_out_buf()
        return (self._out_buf, pyaudio.paContinue)

    def _load_out_buf(self):
        out_buf = np.empty_like(self._out_buf)
        for i in range(self._out_buf.size):
            if (self._env_state == 'a' or self._env_state == 'd'
                    or self._env_state == 's' or self._env_state == 'r'
                    or self._env_state == 'c'):
                out_buf[i] = self._tab_wav[self._phase_acc >> self._PHASE_ACC_SHIFT]
                out_buf[i] = int(round(out_buf[i] * self._amp_ratio))
            else:
                out_buf[i] = 0
            self._phase_acc += self._tuning_freq
            self._phase_acc %= self._PHASE_ACC_MAX
        self._out_buf = out_buf

    # interface functions

    def start(self):
        self._stream = self._pa.open(format=pyaudio.paInt16, channels=1, frames_per_buffer=self._out_buf.size,
                rate=self._SAMPLE_RATE_HZ, input=False, output=True, stream_callback=self._pyaudio_callback)
        self._stream.start_stream()
        self._stream_up = True

    def stop(self):
        if self._stream_up:
            self._stream.stop_stream()
            self._stream_up = False

    def close(self):
        self.stop()
        self._stream.close()
        self._pa.terminate()

    def sel_sin(self):
        self._tab_wav = self._TAB_SIN

    def sel_sqr(self):
        self._tab_wav = self._TAB_SQR

    def sel_saw(self):
        self._tab_wav = self._TAB_SAW

    def sel_rev_saw(self):
        self._tab_wav = self._TAB_REV_SAW

    def sel_tri(self):
        self._tab_wav = self._TAB_TRI

    def note_on(self):
        self._env_state = 'a'

    def note_off(self):
        self._env_state = 'i'

    def set_amp_ratio(self, amp):
        self._amp_ratio = amp

    def set_freq_hz(self, freq):
        self._freq_hz = freq
        self._tuning_freq = int(round(
                (self._freq_hz / self._SAMPLE_RATE_HZ) * self._PHASE_ACC_MAX))

    def set_phase_ratio(self, phase):
        self._phase_acc = int(round(phase * self._PHASE_ACC_MAX)) % self._PHASE_ACC_MAX

    def set_attack_ms(self, attack):
        return

    def set_decay_ms(self, decay):
        return

    def set_sustain_ratio(self, sustain):
        return

    def set_release_ms(self, release):
        return

    def set_cancel_ms(self, cancel):
        return

    def ena_retrigger(self):
        return

    def dis_retrigger(self):
        return

    def ena_recycle(self):
        return

    def dis_recycle(self):
        return

    def set_decimation(self,decimation):
        return

    def dis_decimation(self):
        return

    def set_bit_depth(self,bit_depth):
        return

    def dis_bit_depth(self):
        return

    # wave and envelope lookup tables

    def _load_sin(self):
        self._verbose("Generating sine wave table...")
        self._TAB_SIN = np.empty(self._TAB_LEN, dtype=np.int16)
        for i in range(self._TAB_SIN.size):
            sine_phase = 2 * np.pi * (float(i) / self._TAB_SIN.size)
            self._TAB_SIN[i] = int(round(self._SAMPLE_ABSMAX * np.sin(sine_phase)))
        self._verbose("Sine wave table loaded.")

    def _load_sqr(self):
        self._verbose("Generating square wave table...")
        self._TAB_SQR = np.empty(self._TAB_LEN, dtype=np.int16)
        for i in range(self._TAB_SQR.size):
            if i <= self._TAB_SQR.size / 2:
                self._TAB_SQR[i] = self._SAMPLE_ABSMAX
            else:
                self._TAB_SQR[i] = -1 * self._SAMPLE_ABSMAX
        self._verbose("Square wave table loaded.")

    def _load_saw(self):
        self._verbose("Generating saw wave table...")
        self._TAB_SAW = np.empty(self._TAB_LEN, dtype=np.int16)
        saw_inc = (2 * self._SAMPLE_ABSMAX) / (self._TAB_SAW.size - 1)
        for i in range(self._TAB_SAW.size):
            self._TAB_SAW[i] = int(round((-1
                    * self._SAMPLE_ABSMAX) + i * saw_inc))
        self._verbose("Saw wave table loaded.")

    def _load_rev_saw(self):
        self._verbose("Generating reverse saw wave table...")
        self._TAB_REV_SAW = np.empty(self._TAB_LEN, dtype=np.int16)
        saw_dec = (2 * self._SAMPLE_ABSMAX) / (self._TAB_REV_SAW.size - 1)
        for i in range(self._TAB_REV_SAW.size):
            self._TAB_REV_SAW[i] = int(round(self._SAMPLE_ABSMAX - i * saw_dec))
        self._verbose("Reverse saw wave table loaded.")

    def _load_tri(self):
        self._verbose("Generating triangle wave table...")
        self._TAB_TRI = np.empty(self._TAB_LEN, dtype=np.int16)
        tri_inc = (2 * self._SAMPLE_ABSMAX) / (self._TAB_TRI.size / 2 - 1)
        for i in range(self._TAB_TRI.size // 4):
            self._TAB_TRI[i] = int(round(i * tri_inc))
        for i in range(self._TAB_TRI.size // 2):
            self._TAB_TRI[self._TAB_TRI.size // 4 + i] = int(round(self._SAMPLE_ABSMAX - i * tri_inc))
        for i in range(self._TAB_TRI.size // 4):
            self._TAB_TRI[(3 * self._TAB_TRI.size) // 4 + i] = int(round(-1 * self._SAMPLE_ABSMAX + (i) * tri_inc))
        self._verbose("Triangle wave table loaded.")

    def _load_exp(self):
        self._verbose("Generating exponential ramp table...")
        tab_exp = np.empty(self._TAB_LEN, dtype=np.float64)
        exp_multiplier = self._ENV_VAL_ABSMAX**(1 / tab_exp.size)
        tab_exp[0] = exp_multiplier
        for i in range(1, tab_exp.size):
            tab_exp[i] = tab_exp[i-1] * exp_multiplier
        tab_exp = np.rint(tab_exp)
        self._TAB_EXP = np.empty(tab_exp.size, dtype=np.uint16)
        self._TAB_EXP = tab_exp.astype(np.uint16)
        self._verbose("Exponential ramp table loaded.")

    def _load_all_tabs(self):
        self._verbose("Loading lookup tables...")
        self._load_sin()
        self._load_sqr()
        self._load_saw()
        self._load_rev_saw()
        self._load_tri()
        self._load_exp()
        self._verbose("Lookup tables ready.")

# data array plotting functions

    def plot_tab(self, tab_name):
        if tab_name == "SIN" or tab_name == "SINE":
            tab = self._TAB_SIN
        elif tab_name == "SQR" or tab_name == "SQUARE":
            tab = self._TAB_SQR
        elif tab_name == "SAW":
            tab = self._TAB_SAW
        elif (tab_name == "REV_SAW" or tab_name == "REV SAW"
                  or tab_name == "REVSAW" or tab_name == "REVERSE_SAW"
                  or tab_name == "REVERSE SAW" or tab_name == "REVERSESAW"):
            tab = self._TAB_REV_SAW
        elif tab_name == "TRI" or tab_name == "TRIANGLE":
            tab = self._TAB_TRI
        elif tab_name == "EXP" or tab_name == "EXPONENTIAL":
            tab = self._TAB_EXP
        else:
            return "TableNameError"
        plt.plot(range(tab.size),tab)
        plt.show()

    def plot_buf(self):
        plt.plot(range(self.out_buf.size),out_buf.buffer)
        plt.show(block=False)


"""Beam models and chromaticity corrections."""
from __future__ import annotations

from pathlib import Path
import logging
import hashlib
from methodtools import lru_cache
import astropy.coordinates as apc
import astropy.time as apt
import h5py
import numpy as np
import scipy.interpolate as spi
from tqdm import tqdm
import attr

from . import coordinates as coords
from .loss import ground_loss
from . import sky_models

from edges_io.h5 import HDF5Object

from ..config import config
from .. import const
from edges_cal import (
    FrequencyRange,
    modelling as mdl,
)
from .data import BEAM_PATH
from . import types as tp

logger = logging.getLogger(__name__)

# Reference UTC observation time. At this time, the LST is 0.1666 (00:10 Hrs LST) at the
# EDGES location. NOTE: this is used by default, but can be changed by the user anywhere
# it is used.
REFERENCE_TIME = apt.Time("2014-01-01T09:39:42", location=const.edges_location)

@attr.s
class Beam:
    beam = attr.ib()
    frequency = attr.ib()
    elevation = attr.ib()
    azimuth = attr.ib()
    simulator = attr.ib(default=None, converter=str)
    instrument = attr.ib(default=None, converter=str)
    raw_file = attr.ib(default=None, converter=str)

    @classmethod
    def from_hfss(
        cls,
        path: [str, Path],
        frequency: float,
        linear: bool = True,
        theta_min: float = 0,
        theta_max: float = 180,
        theta_resolution: float = 1,
        phi_min: float = 0,
        phi_max: float = 359,
        phi_resolution: float = 1,
    ) -> Beam:
        """
        Create a Beam object from a HFSS file.

        Parameters
        ----------
        path
            Path to the file. Use :meth:`resolve_file` to get the absolute path
            from a a relative path (starting with ':').
        linear
            Whether the beam values are in linear units (or decibels)
        theta_min, theta_max
            Min/Max of the zenith angle (degrees)
        theta_resolution
            Resolution of the zenith angle.
        phi_min, phi_max
            Min/Max of the azimuth angles (degrees)
        phi_resolution
            The resolution of the azimuth angle.

        Returns
        -------
        beam
            The beam object.
        """
        d = np.genfromtxt(path, skip_header=1, delimiter=",")
        theta = np.arange(theta_min, theta_max + theta_resolution, theta_resolution)
        phi = np.arange(phi_min, phi_max + phi_resolution, phi_resolution)

        beam_map = np.zeros((len(theta), len(phi)))
        for i in range(len(theta)):
            mask = d[:, 1] == theta[i]

            this_phi = d[mask, 0]
            this_power = d[mask, 2]

            this_phi[this_phi < 0] += 360
            this_phi, unique_inds = np.unique(this_phi, axis=0, return_index=True)
            this_power = this_power[unique_inds]

            this_power = this_power[np.argsort(this_phi)]

            if not linear:
                this_power = 10 ** (this_power / 10)

            this_power[np.isnan(this_power)] = 0
            beam_map[i, :] = this_power

        return Beam(
            frequency=np.array([frequency]),
            elevation=theta,
            azimuth=phi,
            beam=beam_map,
            simulator="hfss",
            raw_file=path,
        )

    @classmethod
    def from_wipld(cls, path: tp.PathLike, az_antenna_axis: float = 0) -> Beam:
        """Read a WIPL-D beam."""
        with open(path) as fn:
            file_length = 0
            number_of_frequencies = 0

            flag_columns = False
            frequencies_list = []
            for line in fn:
                file_length += 1
                if line[2] == ">":
                    number_of_frequencies += 1
                    frequencies_list.append(float(line[19:32]))
                elif not flag_columns:
                    line_splitted = line.split()
                    number_of_columns = len(line_splitted)
                    flag_columns = True

            rows_per_frequency = (
                file_length - number_of_frequencies
            ) / number_of_frequencies

            output = np.zeros(
                (
                    int(number_of_frequencies),
                    int(rows_per_frequency),
                    int(number_of_columns),
                )
            )

        frequencies = np.array(frequencies_list)

        with open(path) as fn:
            i = -1
            for line in fn:

                if line[2] == ">":
                    i += 1
                    j = -1
                else:
                    j += 1
                    line_splitted = line.split()
                    line_array = np.array(line_splitted, dtype=float)
                    output[i, j, :] = line_array

        # Re-arrange data
        phi_u = np.unique(output[0, :, 0])
        theta_u = np.unique(output[0, :, 1])
        beam = np.zeros((len(frequencies), len(theta_u), len(phi_u)))

        for i in range(len(frequencies)):
            out_2D = output[i, :, :]

            phi = out_2D[:, 0]
            theta = (
                90 - out_2D[:, 1]
            )  # theta is zero at the zenith, and goes to 180 deg
            gain = out_2D[:, 6]

            theta_u = np.unique(theta)
            it = np.argsort(theta_u)
            theta_a = theta_u[it]

            for j in range(len(theta_a)):
                phi_j = phi[theta == theta_a[j]]
                gain_j = gain[theta == theta_a[j]]

                ip = np.argsort(phi_j)
                gp = gain_j[ip]

                beam[i, j, :] = gp

        # Flip beam from theta to elevation
        beam_maps = beam[:, ::-1, :]

        # Change coordinates from theta/phi, to AZ/EL
        el = np.arange(0, 91)
        az = np.arange(0, 360)

        # Shifting beam relative to true AZ (referenced at due North)
        beam_maps_shifted = cls.shift_beam_maps(az_antenna_axis, beam_maps)

        return Beam(
            frequency=frequencies,
            azimuth=az,
            elevation=el,
            beam=beam_maps_shifted,
            simulator="wipl-d",
            raw_file=path,
        )

    @classmethod
    def from_ideal(cls, delta_f=2, f_low=40, f_high=200, delta_az=1, delta_el=1):
        """Create an ideal beam that is completely unity."""
        freq = np.arange(f_low, f_high, delta_f)
        az = np.arange(0, 360, delta_az)
        el = np.arange(0, 90 + 0.1 * delta_el, delta_el)
        return Beam(
            frequency=freq,
            azimuth=az,
            elevation=el,
            beam=np.ones((len(freq), len(el), len(az))),
            simulator="ideal",
        )

    @classmethod
    def from_feko(cls, path: [str, Path], az_antenna_axis: float = 0) -> Beam:
        """
        Read a FEKO beam file.

        Parameters
        ----------
        filename
            The path to the file.
        az_antenna_axis
            The azimuth of the primary antenna axis, in degrees.

        Returns
        -------
        beam_maps
            A ``(Nfreq, Nel, Naz)`` array giving values of the beam. Note that elevation
            and azimuth are always in 1-degree increments.
        freq
            The frequencies at which the beam is defined.
        """
        filename = Path(path)

        data = np.genfromtxt(str(filename))
        frequency = []
        with open(filename) as fl:
            for line in fl.readlines():
                if line.startswith("#FREQUENCY"):
                    line = line.split(" ")
                    indx = line.index("MHz")
                    frequency.append(float(line[indx - 1]))
        freq = FrequencyRange(np.array(frequency))

        # Loading data and convert to linear representation
        beam_maps = np.zeros((len(frequency), 91, 360))
        for i in range(len(frequency)):
            beam_maps[i] = (10 ** (data[(i * 360) : ((i + 1) * 360), 2::] / 10)).T

        # Shifting beam relative to true AZ (referenced at due North)
        # Due to angle of orientation of excited antenna panels relative to due North
        beam_maps = cls.shift_beam_maps(az_antenna_axis, beam_maps)

        return Beam(
            frequency=freq.freq,
            beam=beam_maps,
            azimuth=np.arange(0, 360),
            elevation=np.arange(0, 91),
            simulator="feko",
            raw_file=path,
        )

    @classmethod
    def from_cst(
        cls,
        file_name_prefix: tp.PathLike,
        f_low: int = 40,
        f_high: int = 100,
        freq_p: int = 61,
        theta_p: float = 181,
        phi_p: float = 361,
        az_antenna_axis: float = 0,
    ) -> Beam:
        """
        Read a CST beam file.

        Parameters
        ----------
        filename
            The path to the file.
        az_antenna_axis
            The azimuth of the primary antenna axis, in degrees.
        f_low, f_high
            lower and higher frequency bounds
        freq_p
            Number of frequency points in the simulation file.
        theta_p
            Number of zenith angle points in the simulation file.
        phi_p
            Number of azimuth points in the simulation file.

        Returns
        -------
        beam
            The beam object.
        """
        phi_t = np.zeros((freq_p, theta_p * phi_p))
        frequency = np.linspace(f_low, f_high, freq_p)
        beam_square = np.zeros((freq_p, theta_p, phi_p))
        res = (f_high - f_low) / (freq_p - 1)

        for i in range(freq_p):
            n = int(i * res + f_low)
            fc_file = open(
                file_name_prefix + "farfield (f=" + str(n) + ") [1].txt", "rb"
            )  #
            for x, line in enumerate(fc_file):
                if x > 1:
                    check = 0
                    for o in range(len(line)):

                        if line[o] != "":
                            check = check + 1
                            if check == 3:
                                phi_t[i][x - 2] = float(line.split()[2])

            fc_file.close()
        for i in range(freq_p):
            for x in range(phi_p):
                beam_square[i, :, x] = phi_t[i, x * theta_p : (x + 1) * theta_p]
            beam_square[i, :, 0 : int(phi_p / 2)] = beam_square[
                i, :, int(phi_p / 2) : phi_p - 1
            ]

        freq = FrequencyRange(np.array(frequency))
        beam_square[:, 91:, :] = 0
        beam_maps = np.flip(beam_square[:, :91, :360], axis=1)
        # Shifting beam relative to true AZ (referenced at due North)
        # Due to angle of orientation of excited antenna panels relative to due North
        beam_maps = cls.shift_beam_maps(az_antenna_axis, beam_maps)

        return Beam(
            frequency=freq.freq,
            beam=beam_maps,
            azimuth=np.arange(0, 360),
            elevation=np.arange(0, 91),
            simulator="cst",
            raw_file=file_name_prefix,
        )

    @classmethod
    def from_feko_raw(
        cls,
        file_name_prefix: tp.PathLike,
        ext: str = "txt",
        f_low: int = 40,
        f_high: int = 100,
        freq_p: int = 61,
        theta_p: float = 181,
        phi_p: float = 361,
        az_antenna_axis: float = 0,
    ) -> Beam:
        """
        Read a FEKO beam file.

        Parameters
        ----------
        filename
            The path to the file.
        az_antenna_axis
            The azimuth of the primary antenna axis, in degrees.
        f_low, f_high
            lower and higher frequency bounds
        freq_p
            Number of frequency points in the simulation file.
        theta_p
            Number of zenith angle points in the simulation file.
        phi_p
            Number of azimuth points in the simulation file.

        Returns
        -------
        beam
            The beam object.
        """
        beam_square = np.zeros((freq_p, theta_p, phi_p))
        frequency = np.linspace(f_low, f_high, freq_p)
        f1 = open(str(file_name_prefix) + "_0-90." + ext)
        f2 = open(str(file_name_prefix) + "_91-180." + ext)
        f3 = open(str(file_name_prefix) + "_181-270." + ext)
        f4 = open(str(file_name_prefix) + "_271-360." + ext)

        z = (
            theta_p * 91 + 10
        )  # ---> change this to no.of theta * no.of phi + No.of header lines

        for index, line in enumerate(f1):
            if index % z == 0:
                co = 0
            if index % z >= 10:
                x = list(map(float, line.split()))
                beam_square[int(index / z), co % theta_p, int(co / theta_p)] = 10 ** (
                    x[8] / 10
                )
                co += 1

        z = theta_p * 90 + 10  #

        for index, line in enumerate(f2):
            if index % z == 0:
                co = 0
            if index % z >= 10:
                x = list(map(float, line.split()))
                beam_square[
                    int(index / z), co % theta_p, int(co / theta_p) + 91
                ] = 10 ** (x[8] / 10)
                co += 1

        for index, line in enumerate(f3):
            if index % z == 0:
                co = 0
            if index % z >= 10:
                x = list(map(float, line.split()))
                beam_square[
                    int(index / z), co % theta_p, int(co / theta_p) + 181
                ] = 10 ** (x[8] / 10)
                co += 1

        for index, line in enumerate(f4):
            if index % z == 0:
                co = 0
            if index % z >= 10:
                x = list(map(float, line.split()))
                beam_square[
                    int(index / z), co % theta_p, int(co / theta_p) + 271
                ] = 10 ** (x[8] / 10)
                co += 1

        freq = FrequencyRange(np.array(frequency))
        beam_square[:, 90:, :] = 0
        beam_maps = np.flip(beam_square[:, :91, :360], axis=1)
        # Shifting beam relative to true AZ (referenced at due North)
        # Due to angle of orientation of excited antenna panels relative to due North
        beam_maps = cls.shift_beam_maps(az_antenna_axis, beam_maps)

        return Beam(
            frequency=freq.freq,
            beam=beam_maps,
            azimuth=np.arange(0, 360),
            elevation=np.arange(0, 91),
            simulator="feko",
            raw_file=file_name_prefix,
        )

    @classmethod
    def get_beam_path(cls, band: str, kind: str | None = None) -> Path:
        """Get a standard path to a beam file."""
        pth = BEAM_PATH / band / "default.txt" if not kind else kind + ".txt"
        if not pth.exists():
            raise FileNotFoundError(f"No beam exists for band={band}.")
        return pth

    def at_freq(
        self,
        freq: np.ndarray,
        model: mdl.Model = mdl.Polynomial(
            n_terms=13, transform=mdl.ScaleTransform(scale=75.0)
        ),
    ) -> Beam:
        """
        Interpolate the beam to a new set of frequencies.

        Parameters
        ----------
        freq
            Frequencies to interpolate to.

        Returns
        -------
        beam
            The Beam object at the new frequencies.
        """
        if len(self.frequency) < 3:
            raise ValueError(
                "Can't freq-interpolate beam that has fewer than three frequencies."
            )

        # Frequency interpolation
        interp_beam = np.zeros((len(freq), len(self.elevation), len(self.azimuth)))
        cached_model = model.at(x=self.frequency)
        for i, bm in enumerate(self.beam.T):
            for j, b in enumerate(bm):
                model_fit = cached_model.fit(ydata=b)
                interp_beam[:, j, i] = model_fit.evaluate(freq)

        return Beam(
            frequency=freq,
            azimuth=self.azimuth,
            elevation=self.elevation,
            beam=interp_beam,
        )

    def smoothed(
        self,
        model: mdl.Model = mdl.Polynomial(n_terms=12),
    ) -> Beam:
        """
        Smoothes the beam within its same set of frequencies.

        ----------

        Returns
        -------
        beam
            The Beam object smoothed over its same frequencies.
        """
        if len(self.frequency) < 3:
            raise ValueError(
                "Can't freq-interpolate beam that has fewer than three frequencies."
            )

        # Frequency smoothing
        smooth_beam = np.zeros_like(self.beam)
        cached_model = model.at(x=self.frequency)
        for i, bm in enumerate(self.beam.T):
            for j, b in enumerate(bm):
                model_fit = cached_model.fit(ydata=b)
                smooth_beam[:, j, i] = model_fit.evaluate()
        return Beam(
            frequency=self.frequency,
            azimuth=self.azimuth,
            elevation=self.elevation,
            beam=smooth_beam,
        )

    @staticmethod
    def shift_beam_maps(az_antenna_axis: float, beam_maps: np.ndarray):
        """Rotate beam maps around an axis.

        Parameters
        ----------
        az_antenna_axis
            The aximuth angle of the antenna axis.
        beam_maps
            Beam maps as a function of frequency, za and az.

        Returns
        -------
        beam maps
            Array of the same shape as the input, but rotated.
        """
        if az_antenna_axis < 0:
            index = -az_antenna_axis
            bm1 = beam_maps[:, :, index::]
            bm2 = beam_maps[:, :, 0:index]
            return np.append(bm1, bm2, axis=2)
        elif az_antenna_axis > 0:
            index = az_antenna_axis
            bm1 = beam_maps[:, :, 0:(-index)]
            bm2 = beam_maps[:, :, (360 - index) : :]
            return np.append(bm2, bm1, axis=2)
        else:
            return beam_maps

    @classmethod
    def resolve_file(
        cls,
        path: tp.PathLike,
        band: str | None = None,
        configuration: str = "default",
        simulator: str = "feko",
    ) -> Path:
        """Resolve a file path to a standard location."""
        if str(path) == ":":
            if band is None:
                raise ValueError("band must be given if path starts with a colon (:)")

            # Get the default beam file.
            return BEAM_PATH / band / f"{configuration}.txt"
        elif str(path).startswith(":"):
            if band is None:
                raise ValueError("band must be given if path starts with a colon (:)")

            # Use a beam file in the standard directory.
            return (
                Path(config["paths"]["beams"]).expanduser()
                / f"{band}/simulations/{simulator}/{str(path)[1:]}"
            ).absolute()
        else:
            return Path(path).absolute()

    @classmethod
    def from_file(
        cls,
        band: str | None,
        simulator: str = "feko",
        beam_file: tp.PathLike = ":",
        configuration: str = "default",
        rotation_from_north: float = 90,
    ) -> Beam:
        """Read a beam from file."""
        beam_file = cls.resolve_file(beam_file, band, configuration, simulator)

        if simulator == "feko":
            rotation_from_north -= 90
            out = cls.from_feko(beam_file, az_antenna_axis=rotation_from_north)
        elif simulator == "wipl-d":  # Beams from WIPL-D
            out = cls.from_wipld(beam_file, rotation_from_north)
        elif simulator == "hfss":
            out = cls.from_hfss(beam_file)
        else:
            raise ValueError(f"Unknown value for simulator: '{simulator}'")

        out.instrument = band
        return out

    @lru_cache(maxsize=1000)
    def angular_interpolator(self, freq_indx: int):
        """Return a callable function that interpolates the beam.

        The returned function has the signature ``interp(az, el)``, where ``az`` is
        azimuth in degrees, and ``el`` is elevation in degrees. They may be arrays,
        in which case they should be the same length.
        """
        el = self.elevation * np.pi / 180 + np.pi / 2

        beam = self.beam[freq_indx]

        if el[-1] > 0.999 * np.pi:
            el = el[:-1]
            beam = beam[:-1]

        spl = spi.RectSphereBivariateSpline(
            el,
            self.azimuth * np.pi / 180,
            beam,
            pole_values=(None, beam[-1]),
            pole_exact=True,
        )
        return lambda az, el: spl(
            el * np.pi / 180 + np.pi / 2, az * np.pi / 180, grid=False
        )

    def between_freqs(self, low=0, high=np.inf) -> Beam:
        """Return a new :class:`Beam` object restricted a given frequency range."""
        mask = (self.frequency >= low) & (self.frequency <= high)
        return attr.evolve(self, frequency=self.frequency[mask], beam=self.beam[mask])

    def get_beam_solid_angle(self) -> float:
        """Calculate the integrated beam solid angle."""
        # Theta vector valid for the FEKO beams. In the beam, dimension 1 increases
        # from index 0 to 90 corresponding to elevation, which is 90-theta
        theta = self.elevation[::-1]

        sin_theta = np.sin(theta * (np.pi / 180))
        sin_theta_2D = np.tile(sin_theta, (360, 1)).T

        beam_integration = np.sum(self.beam * sin_theta_2D)
        return (1 / (4 * np.pi)) * ((np.pi / 180) ** 2) * beam_integration

class BeamFactor(HDF5Object):
    """A non-interpolated beam factor."""

    _structure = {
        "frequency": lambda x: (x.ndim == 1 and x.dtype == float),
        "lst": lambda x: (x.ndim == 1 and x.dtype == float),
        "antenna_temp_above_horizon": lambda x: (x.ndim == 2 and x.dtype == float),
        "loss_fraction": lambda x: (x.ndim == 2 and x.dtype == float),
        "beam_factor": lambda x: (x.ndim == 2 and x.dtype == float),
        "meta": {
            "beam_file": lambda x: isinstance(x, str),
            "simulator": lambda x: isinstance(x, str),
            "f_low": lambda x: isinstance(x, float),
            "f_high": lambda x: isinstance(x, float),
            "normalize_beam": lambda x: isinstance(x, (bool, np.bool_)),
            "sky_model": lambda x: isinstance(x, str),
            "rotation_from_north": lambda x: isinstance(x, float),
            "index_model": lambda x: isinstance(x, str),
            "reference_frequency": lambda x: isinstance(x, float),
            "max_nside": lambda x: isinstance(x, (int, np.int64)),
        },
    }

    @property
    def beam_factor(self):
        """The beam chromaticity factor. Array is ``(n_lsts, n_freq)``."""
        return self.load("beam_factor")

    @property
    def frequency(self):
        """The frequencies at which the beam factor is defined."""
        return self.load("frequency")

    @property
    def lsts(self):
        """The LSTs at which the beam factor is defined."""
        return self.load("lst")

def sky_convolution_generator(
    lsts: np.ndarray,
    ground_loss_file: str,
    beam: Beam,
    sky_model: sky_models.SkyModel,
    index_model: sky_models.IndexModel,
    normalize_beam: bool,
    beam_smoothing: bool,
    smoothing_model: mdl.Model,
    location: apc.EarthLocation = const.edges_location,
    ref_time: apt.Time = REFERENCE_TIME,
):
    """
    Iterate through given LSTs and generate a beam*sky product at each freq and LST.

    This is a generator, so it will yield a single item at a time (to save on memory).

    Parameters
    ----------
    lsts
        The LSTs at which to evaluate the convolution.
    ground_loss_file
        A path to a file containing ground loss information.
    beam
        The beam to convolve.
    sky_model
        The sky model to convolve
    index_model
        The spectral index model of the sky model.
    normalize_beam
        Whether to ensure the beam is properly normalised.
    beam_interpolation
        Whether to smooth over freq axis

    Yields
    ------
    i
        The LST enumerator
    j
        The frequency enumerator
    mean_conv_temp
        The mean temperature after multiplying by the beam (above the horizon)
    conv_temp
        An array containing the temperature after multiuplying by the beam in each pixel
        above the horizon.
    sky
        An array containing the sky temperature in pixel above the horizon.
    beam
        An array containing the interpolatedbeam in pixels above the horizon
    time
        The local time at each LST.
    n_pixels
        The total number of pixels that are not masked.

    Examples
    --------
    Use this function as follows:

    >>> for i, j, mean_t, conv_t, sky, bm, time, npix in sky_convolution_generator():
    >>>     print(conv_t)
    """
    if beam_smoothing:
        beam = beam.smoothed(smoothing_model)
    sky_map = sky_model.at_freq(
        beam.frequency,
        index_model=index_model,
    )

    ground_gain = ground_loss(
        ground_loss_file, band=beam.instrument, freq=beam.frequency
    )
    galactic_coords = sky_model.get_sky_coords()

    # Get the local times corresponding to the given LSTs
    times = coords.lsts_to_times(lsts, ref_time, location)

    for i, time in tqdm(enumerate(times), unit="LST"):
        # Transform Galactic coordinates of Sky Model to Local coordinates
        altaz = galactic_coords.transform_to(
            apc.AltAz(
                location=const.edges_location,
                obstime=time,
            )
        )
        az = np.asarray(altaz.az.deg)
        el = np.asarray(altaz.alt.deg)

        # Number of pixels over 4pi that are not 'nan'
        n_pix_tot = len(el)

        # Selecting coordinates above the horizon
        horizon_mask = el >= 0
        az_above_horizon = az[horizon_mask]
        el_above_horizon = el[horizon_mask]

        # Selecting sky data above the horizon
        sky_above_horizon = np.ones_like(sky_map) * np.nan
        sky_above_horizon[horizon_mask, :] = sky_map[horizon_mask, :]

        # Loop over frequency
        for j in tqdm(range(len(beam.frequency)), unit="Frequency"):
            beam_above_horizon = np.ones(len(sky_map)) * np.nan
            beam_above_horizon[horizon_mask] = beam.angular_interpolator(j)(
                az_above_horizon, el_above_horizon
            )

            n_pix_ok = np.sum(~np.isnan(beam_above_horizon))

            # The nans are only above the horizon
            n_pix_tot_no_nan = n_pix_tot - (len(el_above_horizon) - n_pix_ok)

            if normalize_beam:
                solid_angle = np.nansum(beam_above_horizon) / n_pix_tot_no_nan
                beam_above_horizon *= ground_gain[j] / solid_angle

            antenna_temperature_above_horizon = (
                beam_above_horizon * sky_above_horizon[:, j]
            )

            yield (
                i,
                j,
                np.nansum(antenna_temperature_above_horizon) / n_pix_tot_no_nan,
                antenna_temperature_above_horizon,
                sky_above_horizon,
                beam_above_horizon,
                time,
                n_pix_tot_no_nan,
            )

def simulate_spectra(
    beam: Beam,
    ground_loss_file: [str, Path] = ":",
    f_low: [None, float] = 0,
    f_high: [None, float] = np.inf,
    normalize_beam: bool = True,
    sky_model: sky_models.SkyModel = sky_models.Haslam408(),
    index_model: sky_models.IndexModel = sky_models.ConstantIndex(),
    lsts: np.ndarray = None,
    beam_smoothing: bool = True,
    smoothing_model: mdl.Model = mdl.Polynomial(n_terms=12),
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Simulate global spectra from sky and beam models.

    Parameters
    ----------
    band
        The band of the antenna (low, mid, high).
    beam
        A :class:`Beam` object.
    ground_loss_file
        A file pointing to a ground-loss model. By default, gets the default ground loss
        model for the given ``band``.
    f_low
        Minimum frequency to keep in the simulation (frequencies otherwise defined by
        the beam).
    f_high
        Maximum frequency to keep in the simulation (frequencies otherwise defined by
        the beam).
    normalize_beam
        Whether to normalize the beam to be maximum unity.
    sky_model
        A sky model to use.
    index_model
        An :class:`IndexModel` to use to generate different frequencies of the sky
        model.
    lsts
        The LSTs at which to simulate

    Returns
    -------
    antenna_temperature_above_horizon
        The antenna temperature for pixels above the horizon, shape (Nlst, Nfreq)
    freq
        The frequencies at which the simulation is defined.
    lst
        The LSTs at which the sim is defined.
    """
    beam = beam.between_freqs(f_low, f_high)
    if lsts is None:
        lsts = np.arange(0, 24, 0.5)

    antenna_temperature_above_horizon = np.zeros((len(lsts), len(beam.frequency)))
    for i, j, temperature, _, _, _, _, _ in sky_convolution_generator(
        lsts,
        ground_loss_file,
        beam,
        sky_model,
        index_model,
        normalize_beam,
        beam_smoothing,
        smoothing_model,
    ):
        antenna_temperature_above_horizon[i, j] = temperature

    return antenna_temperature_above_horizon, beam.frequency, lsts

def antenna_beam_factor(
    beam: Beam,
    ground_loss_file: [str, Path] = ":",
    f_low: float = 0,
    f_high: float = np.inf,
    normalize_beam: bool = True,
    sky_model: sky_models.SkyModel = sky_models.Haslam408(),
    index_model: sky_models.IndexModel = sky_models.GaussianIndex(),
    lsts: [None, np.ndarray] = None,
    save_dir: [None, str, Path] = None,
    save_fname: [None, str, Path, bool] = None,
    reference_frequency: [None, float] = None,
    beam_smoothing: bool = True,
    smoothing_model: mdl.Model = mdl.Polynomial(n_terms=12),
):
    """
    Calculate the antenna beam factor.

    Parameters
    ----------
    beam
        A :class:`Beam` object.
    ground_loss_file
        A file pointing to a ground-loss model. By default, gets the default ground loss
        model for the given ``band``.
    f_low
        Minimum frequency to keep in the simulation (frequencies otherwise defined by
        the beam).
    f_high
        Maximum frequency to keep in the simulation (frequencies otherwise defined by
        the beam).
    normalize_beam
        Whether to normalize the beam to be maximum unity.
    sky_model
        A sky model to use.
    index_model
        An :class:`IndexModel` to use to generate different frequencies of the sky
        model.
    twenty_min_per_lst
        How many periods of twenty minutes fit into each LST bin.
    save_dir
        The directory in which to save the output beam factor.
    save_fname
        The filename to save the output beam factor.
    reference_frequency
        The frequency to take as the "reference", i.e. where the chromaticity will
        be by construction unity.

    Returns
    -------
    beam_factor : :class`BeamFactor` instance
    """
    if not save_dir:
        save_dir = Path(config["paths"]["beams"]) / f"{beam.instrument}/beam_factors/"
    else:
        save_dir = Path(save_dir)

    if str(save_fname).startswith(":"):
        save_fname = Path(save_dir).absolute() / str(save_fname)[1:]

    beam = beam.between_freqs(f_low, f_high)

    if lsts is None:
        lsts = np.arange(0, 24, 0.5)

    # Get index of reference frequency
    if reference_frequency is None:
        reference_frequency = (f_high + f_low) / 2

    indx_ref_freq = np.argwhere(beam.frequency >= reference_frequency)[0][0]
    reference_frequency = beam.frequency[indx_ref_freq]

    antenna_temperature_above_horizon = np.zeros((len(lsts), len(beam.frequency)))
    convolution_ref = np.zeros((len(lsts), len(beam.frequency)))
    loss_fraction = np.zeros((len(lsts), len(beam.frequency)))

    for i, j, temperature, _, sky, bm, _, npix_no_nan in sky_convolution_generator(
        lsts,
        beam=beam,
        sky_model=sky_model,
        index_model=index_model,
        normalize_beam=normalize_beam,
        beam_smoothing=beam_smoothing,
        smoothing_model=smoothing_model,
        ground_loss_file=ground_loss_file,
    ):
        antenna_temperature_above_horizon[i, j] = temperature
        # Convolution between (beam at all frequencies) and (sky at reference frequency)
        convolution_ref[i, j] = np.nansum(bm * sky[:, indx_ref_freq])

        # Loss fraction
        loss_fraction[i, j] = 1 - np.nansum(bm) / npix_no_nan

    # Beam factor
    beam_factor = (convolution_ref.T / convolution_ref[:, indx_ref_freq]).T

    out = {
        "frequency": beam.frequency.astype(float),
        "lst": np.array(lsts).astype(float),
        "antenna_temp_above_horizon": antenna_temperature_above_horizon,
        "loss_fraction": loss_fraction,
        "beam_factor": beam_factor,
        "meta": {
            "beam_file": str(beam.raw_file),
            "simulator": beam.simulator,
            "f_low": float(f_low),
            "f_high": float(f_high),
            "normalize_beam": bool(normalize_beam),
            "sky_model": str(sky_model),
            "index_model": str(index_model),
            "reference_frequency": float(reference_frequency),
            "rotation_from_north": float(90),
            "max_nside": int(sky_model.max_res or 0),
        },
    }

    if save_fname is None:
        hsh = hashlib.md5(repr(out["meta"]).encode()).hexdigest()
        save_fname = save_dir / (
            f"{beam.simulator}_{sky_model.__class__.__name__}_"
            f"ref{reference_frequency:.2f}_{hsh}.h5"
        )

    logger.info(f"Writing out beam file to {save_fname}")
    bf = BeamFactor.from_data(out, filename=save_fname)
    bf.write(clobber=True)

    return bf

class InterpolatedBeamFactor(HDF5Object):
    _structure = {
        "beam_factor": None,
        "frequency": None,
        "lst": None,
    }

    @classmethod
    def from_beam_factor(
        cls,
        beam_factor_file,
        band: str | None = None,
        lst_new: np.ndarray | None = None,
        f_new: np.ndarray | None = None,
    ):
        """
        Interpolate beam factor to a new set of LSTs and frequencies.

        The LST interpolation is done using `griddata`, whilst the frequency
        interpolation is done using a polynomial fit.

        Parameters
        ----------
        beam_factor_file : path
            Path to a file containing beam factors produced by
            :func:`antenna_beam_factor`.
            If just a filename (no path), the `beams/band/beam_factors/` directory will
            be searched (dependent on the configured "beams" directory).
        band : str, optional
            If `beam_factor_file` is relative, the band is required to find the file.
        lst_new : array-like, optional
            The LSTs to interpolate to. By default, keep same LSTs as input.
        f_new : array-like, optional
            The frequencies to interpolate to. By default, keep same frequencies as
            input.
        """
        beam_factor_file = Path(beam_factor_file).expanduser()
        if str(beam_factor_file).startswith(":"):
            beam_factor_file = (
                Path(config["paths"]["beams"])
                / band
                / "beam_factors"
                / str(beam_factor_file)[1:]
            )

        if not beam_factor_file.exists():
            raise ValueError(f"The beam factor file {beam_factor_file} does not exist!")

        with h5py.File(beam_factor_file, "r") as fl:
            beam_factor = fl["beam_factor"][...]
            freq = fl["frequency"][...]
            lst = fl["lst"][...]

        # Wrap beam factor and LST for 24-hr interpolation
        beam_factor = np.vstack((beam_factor[-1], beam_factor, beam_factor[0]))
        lst0 = np.append(lst[-1] - 24, lst)
        lst = np.append(lst0, lst[0] + 24)

        if lst_new is not None:
            beam_factor_lst = np.zeros((len(lst_new), len(freq)))
            for i, bf in enumerate(beam_factor.T):
                beam_factor_lst[:, i] = spi.interp1d(lst, bf, kind="cubic")(lst_new)
            lst = lst_new
        else:
            beam_factor_lst = beam_factor

        if f_new is not None:
            # Interpolating beam factor to high frequency resolution
            beam_factor_freq = np.zeros((len(beam_factor_lst), len(f_new)))
            for i, bf in enumerate(beam_factor_lst):
                beam_factor_freq[i] = spi.interp1d(freq, bf, kind="cubic")(f_new)

            freq = f_new
        else:
            beam_factor_freq = beam_factor_lst

        return cls.from_data(
            {"beam_factor": beam_factor_freq, "frequency": freq, "lst": lst}
        )

    def evaluate(self, lst):
        """Fast nearest-neighbour evaluation of the beam factor for a given LST."""
        beam_factor = np.zeros((len(lst), len(self["frequency"])))

        for i, lst_ in enumerate(lst):
            index = np.argmin(np.abs(self["lst"] - lst_))
            beam_factor[i] = self["beam_factor"][index]

        return beam_factor

import numpy,import scipy,import astropy