text:
#
# This file is part of Brazil Data Cube Collection Builder.
# Copyright (C) 2019-2020 INPE.
#
# Brazil Data Cube Collection Builder is free software; you can redistribute it and/or modify it
# under the terms of the MIT License; see LICENSE file for more details.
#

"""Define Brazil Data Cube utils."""

# Python Native
import datetime
import json
import logging
from json import loads as json_parser
from os import remove as resource_remove, path as resource_path
from pathlib import Path
from typing import List
from zipfile import BadZipfile, ZipFile
from zlib import error as zlib_error

# 3rdparty
import boto3
import numpy
import rasterio
import rasterio.features
import rasterio.warp
import requests
import shapely
import shapely.geometry
from bdc_catalog.models import db
from bdc_catalog.utils import multihash_checksum_sha256
from botocore.exceptions import ClientError
from celery import chain, group
from landsatxplore.api import API
from numpngw import write_png
from osgeo import gdal
from rio_cogeo.cogeo import cog_translate
from rio_cogeo.profiles import cog_profiles
from skimage import exposure
from skimage.transform import resize

# Builder
from ..config import CURRENT_DIR, Config
from .sentinel.clients import sentinel_clients

EARTH_EXPLORER_DOWNLOAD_URI = 'https://earthexplorer.usgs.gov/download/{folder}/{sid}/EE'
# List of product id - Earth Explorer Catalog
# TODO: We should use from landsatxplore library but it still not supports 1.5.0.
#  Review it when upgrade to 1.5.0 - https://m2m.cr.usgs.gov/api/docs/json/
EARTH_EXPLORER_PRODUCT_ID = {
    'LANDSAT_TM_C1': '5e83d08ee1e2890b',   # Landsat 4/5
    'LANDSAT_ETM_C1': '5e83a506f72553b6',  # Landsat 7
    'LANDSAT_8_C1': '5e83d0b84df8d8c2'     # Landsat 8
}

def get_or_create_model(model_class, defaults=None, engine=None, **restrictions):
    """Get or create Brazil Data Cube model.

    Utility method for looking up an object with the given restrictions, creating one if necessary.

    Args:
        model_class (BaseModel) - Base Model of Brazil Data Cube DB
        defaults (dict) - Values to fill out model instance
        restrictions (dict) - Query Restrictions
    Returns:
        BaseModel Retrieves model instance
    """
    if not engine:
        engine = db

    instance = engine.session.query(model_class).filter_by(**restrictions).first()

    if instance:
        return instance, False

    params = dict((k, v) for k, v in restrictions.items())

    params.update(defaults or {})
    instance = model_class(**params)

    engine.session.add(instance)

    return instance, True

def dispatch(activity: dict, skip_l1=None, **kwargs):
    """Dispatches the activity to the respective celery task handler.

    Args:
        activity (RadcorActivity) - A not done activity
        skip_l1 (bool) - Skip publish schedule for download tasks.
        kwargs (dict) - Extra arguments to pass celery function.
    """
    from .landsat import tasks as landsat_tasks
    from .landsat.utils import factory as landsat_factory
    from .sentinel import tasks as sentinel_tasks
    from .sentinel.utils import factory as sentinel_factory
    # TODO: Implement it as factory (TaskDispatcher) and pass the responsibility to the task type handler

    app = activity.get('activity_type')

    if app == 'downloadS2':
        # Atm Correction chain
        atm_corr_publish_chain = sentinel_tasks.atm_correction.s() | sentinel_tasks.publish_sentinel.s(skip_l1=skip_l1)
        # Publish ATM Correction
        upload_chain = sentinel_tasks.apply_post_processing.s()

        inner_group = upload_chain

        if activity['args'].get('harmonize'):
            # Harmonization chain
            harmonize_chain = sentinel_tasks.harmonization_sentinel.s() | \
                              sentinel_tasks.publish_sentinel.s(skip_l1=skip_l1) | \
                              sentinel_tasks.upload_sentinel.s()
            inner_group = group(upload_chain, harmonize_chain)

        inner_group = atm_corr_publish_chain | inner_group

        after_download_group = [
            inner_group
        ]

        outer_group = group(*after_download_group)
        task_chain = sentinel_tasks.download_sentinel.s(activity) | outer_group
        return chain(task_chain).apply_async()
    elif app == 'correctionS2':
        task_chain = sentinel_tasks.atm_correction.s(activity) | \
                        sentinel_tasks.publish_sentinel.s() | \
                        sentinel_tasks.apply_post_processing.s()
        return chain(task_chain).apply_async()
    elif app == 'publishS2':
        tasks = [sentinel_tasks.publish_sentinel.s(activity), sentinel_tasks.apply_post_processing.s()]

        return chain(*tasks).apply_async()
    elif app == 'harmonizeS2':
        task_chain = sentinel_tasks.harmonization_sentinel.s(activity) | sentinel_tasks.publish_sentinel.s() | \
                    sentinel_tasks.upload_sentinel.s()
        return chain(task_chain).apply_async()
    elif app == 'uploadS2':
        return sentinel_tasks.upload_sentinel.s(activity).apply_async()

    elif app == 'downloadLC8':
        # Atm Correction chain
        atm_corr_chain = landsat_tasks.atm_correction_landsat.s()
        # Publish ATM Correction
        publish_atm_chain = landsat_tasks.publish_landsat.s() | landsat_tasks.upload_landsat.s()

        inner_group = publish_atm_chain

        # Check if will add harmonization chain on group
        if activity['args'].get('harmonize'):
            # Harmonization chain
            harmonize_chain = landsat_tasks.harmonization_landsat.s() | landsat_tasks.publish_landsat.s() | \
                        landsat_tasks.upload_landsat.s()
            inner_group = group(publish_atm_chain, harmonize_chain)

        atm_chain = atm_corr_chain | inner_group
        outer_group = group(atm_chain)
        task_chain = landsat_tasks.download_landsat.s(activity) | outer_group

        return chain(task_chain).apply_async()
    elif app == 'correctionLC8':
        # Atm Correction chain
        atm_corr_chain = landsat_tasks.atm_correction_landsat.s(activity)
        # Publish ATM Correction
        publish_atm_chain = landsat_tasks.publish_landsat.s() | landsat_tasks.upload_landsat.s()

        inner_group = publish_atm_chain

        # Check if will add harmonization chain on group
        if activity['args'].get('harmonize'):
            # Harmonization chain
            harmonize_chain = landsat_tasks.harmonization_landsat.s() | landsat_tasks.publish_landsat.s() | \
                        landsat_tasks.upload_landsat.s()
            inner_group = group(publish_atm_chain, harmonize_chain)

        task_chain = atm_corr_chain | inner_group
        return chain(task_chain).apply_async()
    elif app == 'publishLC8':
        task_chain = landsat_tasks.publish_landsat.s(activity)

        # When collection L2, chain the upload task
        if activity['collection_id'] in list(landsat_factory.map['l2'].keys()):
            task_chain |= landsat_tasks.upload_landsat.s()

        return chain(task_chain).apply_async()
    elif app == 'harmonizeLC8':
        task_chain = landsat_tasks.harmonization_landsat.s(activity) | landsat_tasks.publish_landsat.s() | \
                    landsat_tasks.upload_landsat.s()
        return chain(task_chain).apply_async()
    elif app == 'uploadLC8':
        return landsat_tasks.upload_landsat.s(activity).apply_async()

def create_wkt(ullon, ullat, lrlon, lrlat):
    """Create WKT representation using lat/long coordinates.

    Args:
        ullon - Upper Left longitude
        ullat - Upper Left Latitude
        lrlon - Lower Right Longitude
        lrlat - Lower Right Latitude

    Returns:
        WKT Object from osr
    """
    from ogr import Geometry, wkbLinearRing, wkbPolygon

    # Create ring
    ring = Geometry(wkbLinearRing)
    ring.AddPoint(ullon, ullat)
    ring.AddPoint(lrlon, ullat)
    ring.AddPoint(lrlon, lrlat)
    ring.AddPoint(ullon, lrlat)
    ring.AddPoint(ullon, ullat)

    # Create polygon
    poly = Geometry(wkbPolygon)
    poly.AddGeometry(ring)

    return poly.ExportToWkt(),poly

def get_earth_explorer_api():
    credentials = get_credentials()['landsat']

    return API(credentials['username'], credentials['password'])

def get_landsat_scenes(wlon, nlat, elon, slat, startdate, enddate, cloud, formal_name: str):
    """List landsat scenes from USGS."""
    api = get_earth_explorer_api()

    landsat_folder_id = EARTH_EXPLORER_PRODUCT_ID.get(formal_name)

    if landsat_folder_id is None:
        raise ValueError('Invalid Landsat product name. Expected one of {}'.format(EARTH_EXPLORER_PRODUCT_ID.keys()))

    # Request
    scenes_result = api.search(
        dataset=formal_name,
        bbox=(slat, wlon, nlat, elon),  # TODO: Pass xmin,ymin,xmax,ymax when landsatxplore bug are fixed.
        start_date=startdate,
        end_date=enddate,
        max_cloud_cover=cloud or 100,
        max_results=50000
    )

    scenes_output = {}

    for scene in scenes_result:
        if scene['displayId'].endswith('RT') or scene['displayId'].startswith('LO08'):
            logging.warning('Skipping RealTime/OLI-Only {}'.format(scene['displayId']))
            continue

        copy_scene = dict()
        copy_scene['sceneid'] = scene['displayId']
        copy_scene['scene_id'] = scene['entityId']
        copy_scene['cloud'] = float(scene['cloudCover'])
        copy_scene['date'] = scene['acquisitionDate']

        xmin, ymin, xmax, ymax = [float(value) for value in scene['sceneBounds'].split(',')]

        # TODO: Check data integrity
        # Sometimes the USGS responds invalid bounding box scenes while searching in EarthExplorer Catalog.
        # w=-60.87065, n=-10.18204, e=-57.66829, s=-12.18696
        # The expected scenes are:
        # 228067, 228068, 228069, 229067, 229068, 229069, 230067, 230068, 230069.
        # However, an invalid scene will be found (074068, 074067).
        if xmin - xmax < -3:
            logging.warning(f'Scene {scene["displayId"]} inconsistent.')
            continue

        copy_scene['wlon'] = xmin
        copy_scene['slat'] = ymin
        copy_scene['elon'] = xmax
        copy_scene['nlat'] = ymax
        copy_scene['link'] = EARTH_EXPLORER_DOWNLOAD_URI.format(folder=landsat_folder_id, sid=scene['entityId'])

        pathrow = scene['displayId'].split('_')[2]

        copy_scene['path'] = pathrow[:3]
        copy_scene['row'] = pathrow[3:]

        scenes_output[scene['displayId']] = copy_scene

    return scenes_output

def get_sentinel_scenes(wlon,nlat,elon,slat,startdate,enddate,cloud,limit,productType=None):
    """Retrieve Sentinel Images from Copernicus."""
    #    api_hub options:
    #    'https://scihub.copernicus.eu/apihub/' for fast access to recently acquired imagery in the API HUB rolling archive
    #    'https://scihub.copernicus.eu/dhus/' for slower access to the full archive of all acquired imagery
    scenes = {}
    totres = 1000000
    first = 0
    pquery = 'https://scihub.copernicus.eu/dhus/search?format=json'
    pquery = 'https://scihub.copernicus.eu/apihub/search?format=json'
    pquery += '&q=platformname:Sentinel-2'
    if productType is not None:
        pquery += ' AND producttype:{}'.format(productType)
    if enddate is None:
        enddate = datetime.datetime.now().strftime("%Y-%m-%d")
    pquery += ' AND beginposition:[{}T00:00:00.000Z TO {}T23:59:59.999Z]'.format(startdate,enddate)
    pquery += ' AND cloudcoverpercentage:[0 TO {}]'.format(cloud)
    if wlon == elon and slat == nlat:
        pfootprintWkt,footprintPoly = create_wkt(wlon-0.01,nlat+0.01,elon+0.01,slat-0.01)
        pquery += ' AND (footprint:"Contains({})")'.format(footprintPoly)
    else:
        pfootprintWkt,footprintPoly = create_wkt(wlon,nlat,elon,slat)
        pquery += ' AND (footprint:"Intersects({})")'.format(footprintPoly)

    limit = int(limit)
    rows = min(100,limit)
    count_results = 0

    users = sentinel_clients.users

    if not users:
        raise ValueError('No sentinel user set')

    username = list(users)[0]
    password = users[username]['password']

    while count_results < min(limit,totres) and totres != 0:
        rows = min(100,limit-len(scenes),totres)
        first = count_results
        query = pquery + '&rows={}&start={}'.format(rows,first)
        try:
            # Using sentinel user and release on out of scope
            r = requests.get(query, auth=(username, password), verify=True)

            if not r.status_code // 100 == 2:
                logging.exception('openSearchS2SAFE API returned unexpected response {}:'.format(r.status_code))
                return {}
            r_dict = r.json()
            #logging.warning('r_dict - {}'.format(json.dumps(r_dict, indent=2)))

        except requests.exceptions.RequestException as exc:
            return {}

        if 'entry' in r_dict['feed']:
            totres = int(r_dict['feed']['opensearch:totalResults'])
            logging.warning('Results for this feed: {}'.format(totres))
            results = r_dict['feed']['entry']
            #logging.warning('Results: {}'.format(results))
            if not isinstance(results, list):
                results = [results]
            for result in results:
                count_results += 1
                identifier = result['title']
                type = identifier.split('_')[1]
                ### Jump level 2 images (will download and process only L1C)
                if type == 'MSIL2A':
                    logging.warning('openSearchS2SAFE skipping {}'.format(identifier))
                    continue
                scenes[identifier] = {}
                scenes[identifier]['pathrow'] = identifier.split('_')[-2][1:]
                scenes[identifier]['sceneid'] = identifier
                scenes[identifier]['type'] = identifier.split('_')[1]
                for data in result['date']:
                    if str(data['name']) == 'beginposition':
                        scenes[identifier]['date'] = str(data['content'])[0:10]
                if not isinstance(result['double'], list):
                    result['double'] = [result['double']]
                for data in result['double']:
                    if str(data['name']) == 'cloudcoverpercentage':
                        scenes[identifier]['cloud'] = float(data['content'])
                for data in result['str']:
                    if str(data['name']) == 'size':
                        scenes[identifier]['size'] = data['content']
                    if str(data['name']) == 'footprint':
                        scenes[identifier]['footprint'] = data['content']
                    if str(data['name']) == 'tileid':
                        scenes[identifier]['tileid'] = data['content']
                if 'tileid' not in scenes[identifier]:
                    logging.warning( 'openSearchS2SAFE identifier - {} - tileid {} was not found'.format(identifier,scenes[identifier]['pathrow']))
                    logging.warning(json.dumps(scenes[identifier], indent=4))
                scenes[identifier]['link'] = result['link'][0]['href']
                scenes[identifier]['icon'] = result['link'][2]['href']
        else:
            totres = 0
    return scenes

def is_valid_tif(input_data_set_path):
    """Validate Tif.

    Args:
        input_data_set_path (str) - Path to the input data set
    Returns:
        True if tif is valid, False otherwise
    """
    # this allows GDAL to throw Python Exceptions
    gdal.UseExceptions()

    try:
        ds = gdal.Open(str(input_data_set_path))
        srcband = ds.GetRasterBand(1)

        array = srcband.ReadAsArray()

        # Check if min == max
        if array.min() == array.max() == 0:
            del ds
            return False
        del ds
        return True
    except RuntimeError as e:
        logging.error('Unable to open {} {}'.format(input_data_set_path, e))
        return False

def load_img(img_path):
    """Load an image."""
    try:
        with rasterio.open(img_path) as dataset:
            img = dataset.read(1).flatten()
        return img
    except:
        logging.error('Cannot find {}'.format(img_path))
        raise RuntimeError('Cannot find {}'.format(img_path))

def extractall(file, destination=None):
    """Extract zipfile."""
    archive = ZipFile(file, 'r')

    if destination is None:
        destination = resource_path.dirname(file)

    archive.extractall(destination)
    archive.close()

def get_credentials():
    """Retrieve global secrets with credentials."""
    file = resource_path.join(resource_path.dirname(CURRENT_DIR), 'secrets.json')

    with open(file) as f:
        return json_parser(f.read())

def generate_cogs(input_data_set_path, file_path, profile='lzw', profile_options=None, **options):
    """Generate Cloud Optimized GeoTIFF files (COG).

    Example:
        >>> from bdc_collection_builder.collections.utils import generate_cogs
        >>> import gdal
        >>>
        >>> tif_file = '/path/to/tif'
        >>> generate_cogs(tif_file, '/tmp/cog.tif')

    Args:
        input_data_set_path (str) - Path to the input data set
        file_path (str) - Target data set filename
        profile (str) - A COG profile based in `rio_cogeo.profiles`.
        profile_options (dict) - Custom options to the profile.

    Returns:
        Path to COG.
    """
    if profile_options is None:
        profile_options = dict()

    output_profile = cog_profiles.get(profile)
    output_profile.update(dict(BIGTIFF="IF_SAFER"))
    output_profile.update(profile_options)

    # Dataset Open option (see gdalwarp `-oo` option)
    config = dict(
        GDAL_NUM_THREADS="ALL_CPUS",
        GDAL_TIFF_INTERNAL_MASK=True,
        GDAL_TIFF_OVR_BLOCKSIZE="128",
    )

    cog_translate(
        str(input_data_set_path),
        str(file_path),
        output_profile,
        config=config,
        in_memory=False,
        quiet=True,
        **options,
    )
    return str(file_path)

def is_valid_compressed(file):
    """Check tar gz or zip is valid."""
    try:
        archive = ZipFile(file, 'r')
        try:
            corrupt = archive.testzip()
        except zlib_error:
            corrupt = True
        archive.close()
    except BadZipfile:
        corrupt = True

    return not corrupt

def extract_and_get_internal_name(zip_file_name, extract_to=None):
    """Extract zipfile and return internal folder path."""
    # Check if file is valid
    valid = is_valid_compressed(zip_file_name)

    if not valid:
        raise IOError('Invalid zip file "{}"'.format(zip_file_name))
    else:
        extractall(zip_file_name, destination=extract_to)

        # Get extracted zip folder name
        with ZipFile(zip_file_name) as zipObj:
            listOfiles = zipObj.namelist()
            extracted_file_path = listOfiles[0]

        return extracted_file_path

def upload_file(file_name, bucket='bdc-ds-datacube', object_name=None):
    """Upload a file to an S3 bucket.

    Adapted code from boto3 example.

    Args:
        file_name (str|_io.TextIO): File to upload
        bucket (str): Bucket to upload to
        object_name (str): S3 object name. If not specified then file_name is used
    """
    # If S3 object_name was not specified, use file_name
    if object_name is None:
        object_name = file_name

    # Upload the file
    s3_client = boto3.client('s3', region_name=Config.AWS_REGION_NAME, aws_access_key_id=Config.AWS_ACCESS_KEY_ID, aws_secret_access_key=Config.AWS_SECRET_ACCESS_KEY)
    try:
        s3_client.upload_file(file_name, bucket, object_name)
    except ClientError as e:
        logging.error(e)
        return False
    return True

def generate_evi_ndvi(red_band: str, nir_band: str, blue_bland: str, evi_name: str, ndvi_name: str):
    """Generate Normalized Difference Vegetation Index (NDVI) and Enhanced Vegetation Index (EVI).

    Args:
        red_band: Path to the RED band
        nir_band: Path to the NIR band
        blue_bland: Path to the BLUE band
        evi_name: Path to save EVI file
        ndvi_name: Path to save NDVI file
    """
    data_set = gdal.Open(red_band, gdal.GA_ReadOnly)
    raster_xsize = data_set.RasterXSize
    raster_ysize = data_set.RasterYSize
    red = data_set.GetRasterBand(1).ReadAsArray(0, 0, data_set.RasterXSize, data_set.RasterYSize).astype(numpy.float32)/10000.

    # Close data_set
    del data_set

    data_set = gdal.Open(nir_band, gdal.GA_ReadOnly)
    nir = data_set.GetRasterBand(1).ReadAsArray(0, 0, data_set.RasterXSize, data_set.RasterYSize).astype(numpy.float32)/10000.
    nir = resize(nir, red.shape, order=1, preserve_range=True).astype(numpy.float32)

    del data_set
    data_set = gdal.Open(blue_bland, gdal.GA_ReadOnly)
    blue = data_set.GetRasterBand(1).ReadAsArray(0, 0, data_set.RasterXSize, data_set.RasterYSize).astype(numpy.float32)/10000.

    # Create the ndvi image data_set
    remove_file(ndvi_name)

    driver = gdal.GetDriverByName('GTiff')

    raster_ndvi = (10000 * (nir - red) / (nir + red + 0.0001)).astype(numpy.int16)
    ndvi_data_set = driver.Create(ndvi_name, raster_xsize, raster_ysize, 1, gdal.GDT_Int16, options=['COMPRESS=LZW',
                                                                                'TILED=YES'])
    ndvi_data_set.SetGeoTransform(data_set.GetGeoTransform())
    ndvi_data_set.SetProjection(data_set.GetProjection())
    ndvi_data_set.GetRasterBand(1).WriteArray(raster_ndvi)
    del ndvi_data_set

    # Create the evi image data set
    remove_file(evi_name)
    evi_data_set = driver.Create(evi_name, raster_xsize, raster_ysize, 1, gdal.GDT_Int16, options=['COMPRESS=LZW',
                                                                                'TILED=YES'])
    raster_evi = (10000 * 2.5 * (nir - red)/(nir + 6. * red - 7.5 * blue + 1)).astype(numpy.int16)
    evi_data_set.SetGeoTransform(data_set.GetGeoTransform())
    evi_data_set.SetProjection(data_set.GetProjection())
    evi_data_set.GetRasterBand(1).WriteArray(raster_evi)
    del raster_evi
    del evi_data_set

    del data_set

def remove_file(file_path: str):
    """Remove file if exists.

    Throws Error when user doesn't have access to the file at given path
    """
    if resource_path.exists(file_path):
        resource_remove(file_path)

def refresh_assets_view(refresh_on_aws=True):
    """Update the Brazil Data Cube Assets View."""
    if not Config.ENABLE_REFRESH_VIEW:
        logging.info('Skipping refresh view.')
        return

    logging.info('View refreshed.')

def create_quick_look(png_file: str, files: List[str], rows=768, cols=768):
    """Generate a Quick Look file (PNG based) from a list of files.

    Note:
        The file order in ``files`` represents the bands Red, Green and Blue, respectively.

    Exceptions:
        RasterIOError when could not open a raster file band

    Args:
        png_file: Path to store the quicklook file.
        files: List of file paths to open and read the Raster files.
        rows: Image height. Default is 768.
        cols: Image width. Default is 768.
    """
    image = numpy.zeros((rows, cols, len(files),), dtype=numpy.uint8)

    nb = 0
    for band in files:
        with rasterio.open(band) as data_set:
            raster = data_set.read(1)

        raster = resize(raster, (rows, cols), order=1, preserve_range=True)
        nodata = raster == -9999
        # Evaluate minimum and maximum values
        a = numpy.array(raster.flatten())
        p1, p99 = numpy.percentile(a[a>0], (1, 99))
        # Convert minimum and maximum values to 1,255 - 0 is nodata
        raster = exposure.rescale_intensity(raster, in_range=(p1, p99),out_range=(1, 255)).astype(numpy.uint8)
        image[:, :, nb] = raster.astype(numpy.uint8) * numpy.invert(nodata)
        nb += 1

    write_png(str(png_file), image, transparent=(0, 0, 0))

def create_asset_definition(href: str, mime_type: str, role: List[str], absolute_path: str,
                            created=None, is_raster=False):
    """Create a valid asset definition for collections.

    TODO: Generate the asset for `Item` field with all bands

    Args:
        href - Relative path to the asset
        mime_type - Asset Mime type str
        role - Asset role. Available values are: ['data'], ['thumbnail']
        absolute_path - Absolute path to the asset. Required to generate check_sum
        created - Date time str of asset. When not set, use current timestamp.
        is_raster - Flag to identify raster. When set, `raster_size` and `chunk_size` will be set to the asset.
    """
    fmt = '%Y-%m-%dT%H:%M:%S'
    _now_str = datetime.datetime.utcnow().strftime(fmt)

    if created is None:
        created = _now_str
    elif isinstance(created, datetime.datetime):
        created = created.strftime(fmt)

    asset = {
        'href': str(href),
        'type': mime_type,
        'bdc:size': Path(absolute_path).stat().st_size,
        'checksum:multihash': multihash_checksum_sha256(str(absolute_path)),
        'roles': role,
        'created': created,
        'updated': _now_str
    }

    if is_raster:
        with rasterio.open(str(absolute_path)) as data_set:
            asset['bdc:raster_size'] = dict(
                x=data_set.shape[1],
                y=data_set.shape[0],
            )

            chunk_x, chunk_y = data_set.profile.get('blockxsize'), data_set.profile.get('blockxsize')

            if chunk_x is None or chunk_x is None:
                raise RuntimeError('Can\'t compute raster chunk size. Is it a tiled/ valid Cloud Optimized GeoTIFF?')

            asset['bdc:chunk_size'] = dict(x=chunk_x, y=chunk_y)

    return asset

def raster_extent(file_path: str, epsg = 'EPSG:4326') -> shapely.geometry.Polygon:
    """Get raster extent in arbitrary CRS

    Args:
        file_path (str): Path to image
        epsg (str): EPSG Code of result crs

    Returns:
        dict: geojson-like geometry
    """

    with rasterio.open(str(file_path)) as data_set:
        _geom = shapely.geometry.mapping(shapely.geometry.box(*data_set.bounds))
        return shapely.geometry.shape(rasterio.warp.transform_geom(data_set.crs, epsg, _geom, precision=6))

def raster_convexhull(file_path: str, epsg='EPSG:4326') -> dict:
    """get image footprint

    Args:
        file_path (str): image file
        epsg (str): geometry EPSG

    See:
        https://rasterio.readthedocs.io/en/latest/topics/masks.html
    """
    with rasterio.open(str(file_path)) as data_set:
        # Read raster data, masking nodata values
        data = data_set.read(1, masked=True)
        data[data != numpy.ma.masked] = 1
        data[data == numpy.ma.masked] = 0
        data = data.astype(numpy.uint8)
        # Create mask, which 1 represents valid data and 0 nodata
        geoms = []
        res = {'val': []}
        for geom, val in rasterio.features.shapes(data, mask=data, transform=data_set.transform):
            geom = rasterio.warp.transform_geom(data_set.crs, epsg, geom, precision=6)

            res['val'].append(val)
            geoms.append(shapely.geometry.shape(geom))

        if len(geoms) == 1:
            return geoms[0]

        multi_polygons = shapely.geometry.MultiPolygon(geoms)

        return multi_polygons.convex_hull

max_stars_repo_name: betonr/bdc-collection-builder

max_issues_repo_name: betonr/bdc-collection-builder

max_forks_repo_name: betonr/bdc-collection-builder


text:
#!/usr/bin/env python3
#
# This example demonstrates the use of consistent radial transport
# on f_re and n_re.
#
# Run as
#
#   $ ./generate.py
#   $ ../../build/iface/dreami dream_settings.h5
#
# ###################################################################

import numpy as np
import sys

sys.path.append('../../py/')

from DREAM.DREAMSettings import DREAMSettings
import DREAM.Settings.Equations.ColdElectronTemperature as T_cold
import DREAM.Settings.Equations.ElectricField as Efield
import DREAM.Settings.Equations.HotElectronDistribution as FHot
import DREAM.Settings.Equations.IonSpecies as Ions
import DREAM.Settings.Equations.RunawayElectrons as Runaways
import DREAM.Settings.Solver as Solver
import DREAM.Settings.TransportSettings as Transport

ds = DREAMSettings()

# Physical parameters
E = 0.5     # Electric field strength (V/m)
n = 5e19    # Electron density (m^-3)
T = 1e3     # Temperature (eV)

# Grid parameters
pMax = 50     # maximum momentum in units of m_e*c
Np   = 50     # number of momentum grid points
Nxi  = 24     # number of pitch grid points
tMax = 1      # simulation time in seconds
Nt   = 50     # number of time steps
Nr   = 8      # number of radial grid points

dBOverB = 1e-3  # Magnetic perturbation strength
R0 = 1.6      # Tokamak major radius

# Set E_field
#ds.eqsys.E_field.setPrescribedData(E)
ds.eqsys.E_field.setType(Efield.TYPE_SELFCONSISTENT)
ds.eqsys.E_field.setBoundaryCondition(bctype=Efield.BC_TYPE_PRESCRIBED, V_loop_wall_R0=E*2*np.pi, R0=R0)

# Set temperature
ds.eqsys.T_cold.setPrescribedData(T)

# Set ions
ds.eqsys.n_i.addIon(name='D', Z=1, iontype=Ions.IONS_PRESCRIBED_FULLY_IONIZED, n=n)

# Disable avalanche generation
ds.eqsys.n_re.setAvalanche(Runaways.AVALANCHE_MODE_FLUID_HESSLOW)
ds.eqsys.n_re.setDreicer(Runaways.DREICER_RATE_NEURAL_NETWORK)

# Disable hot-tail grid
ds.hottailgrid.setEnabled(False)
# Runaway grid
ds.runawaygrid.setNxi(Nxi)
ds.runawaygrid.setNp(Np)
ds.runawaygrid.setPmax(pMax)

# Set initial hot electron Maxwellian
#ds.eqsys.f_re.setInitialValue(0)

# Set up radial grid
ds.radialgrid.setB0(3)
ds.radialgrid.setMinorRadius(0.5)
ds.radialgrid.setWallRadius(0.5)
ds.radialgrid.setNr(Nr)

# Set Rechester-Rosenbluth transport
ds.eqsys.f_re.transport.setMagneticPerturbation(dBOverB)
ds.eqsys.f_re.transport.setBoundaryCondition(Transport.BC_F_0)

# Set solver type
#ds.solver.setType(Solver.LINEAR_IMPLICIT) # semi-implicit time stepping
ds.solver.setType(Solver.NONLINEAR)
ds.solver.setVerbose(False)
ds.solver.setLinearSolver(Solver.LINEAR_SOLVER_MKL)
#ds.solver.tolerance.set(reltol=1e-4)

# include otherquantities to save to output
ds.other.include('fluid')

# Set time stepper
ds.timestep.setTmax(tMax)
ds.timestep.setNt(Nt)

ds.output.setTiming(stdout=True, file=True)
ds.output.setFilename('output.h5')

# Save settings to HDF5 file
ds.save('dream_settings.h5')

ext: py

max_stars_repo_name: chalmersplasmatheory/DREAM

max_issues_repo_name: chalmersplasmatheory/DREAM

max_forks_repo_name: chalmersplasmatheory/DREAM

reason: import numpy


text:
# python3
# Copyright 2018 DeepMind Technologies Limited. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Rescaling layers (e.g. to match action specs)."""

from acme import specs
import haiku as hk
from jax import lax
import jax.numpy as jnp

class ClipToSpec(hk.Module):
  """Sonnet module clipping inputs to within a BoundedArraySpec."""

  def __init__(self, spec: specs.BoundedArray, name: str = 'clip_to_spec'):
    super().__init__(name=name)
    self._min = spec.minimum
    self._max = spec.maximum

  def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
    return jnp.clip(inputs, self._min, self._max)

class RescaleToSpec(hk.Module):
  """Sonnet module rescaling inputs in [-1, 1] to match a BoundedArraySpec."""

  def __init__(self, spec: specs.BoundedArray, name: str = 'rescale_to_spec'):
    super().__init__(name=name)
    self._scale = spec.maximum - spec.minimum
    self._offset = spec.minimum

  def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
    inputs = 0.5 * (inputs + 1.0)  # [0, 1]
    output = inputs * self._scale + self._offset  # [minimum, maximum]

    return output

class TanhToSpec(hk.Module):
  """Sonnet module squashing real-valued inputs to match a BoundedArraySpec."""

  def __init__(self, spec: specs.BoundedArray, name: str = 'tanh_to_spec'):
    super().__init__(name=name)
    self._scale = spec.maximum - spec.minimum
    self._offset = spec.minimum

  def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
    inputs = lax.tanh(inputs)  # [-1, 1]
    inputs = 0.5 * (inputs + 1.0)  # [0, 1]
    output = inputs * self._scale + self._offset  # [minimum, maximum]
    return output

max_stars_repo_name: novatig/acme

max_issues_repo_name: novatig/acme

max_forks_repo_name: novatig/acme

reason: import jax,from jax


text:
#!/usr/bin/env python3
import pybullet as p
import time
import math
from datetime import datetime
from time import sleep
import numpy as np
import random
import cv2
import os
import argparse
import torch

import sys

sys.path.append('./')

from env import Engine
from utils_env import get_view,safe_path,cut_frame,point2traj,get_gripper_pos,backup_code

#################################
import signal
import importlib
import torch
import torch.nn as nn

import sh
import re
import torch.nn.functional as F

np.set_printoptions(precision=4,suppress=True,linewidth=300)

class Engine93(Engine):
    def __init__(self, worker_id, opti, p_id, taskId=5, maxSteps=15, n_dmps=3, cReward=True):
        super(Engine93,self).__init__(opti, wid=worker_id, p_id=p_id, maxSteps=maxSteps, taskId=taskId, n_dmps=n_dmps, cReward=cReward,robot_model=None)
        self.opti = opti
        self.wid = worker_id
        self.p = p_id
        self.p.setPhysicsEngineParameter(enableConeFriction=1)
        self.p.setPhysicsEngineParameter(contactBreakingThreshold=0.001)
        self.p.setPhysicsEngineParameter(allowedCcdPenetration=0.0)

        self.p.setPhysicsEngineParameter(numSolverIterations=20)
        self.p.setPhysicsEngineParameter(numSubSteps=10)

        self.p.setPhysicsEngineParameter(enableFileCaching=0)

        self.p.setTimeStep(1 / 30.0)
        self.p.setGravity(0, 0, -9.81)

        expert_traj_dir = os.path.join(self.robot_recordings_dir,"87-0")
        self.data_q = np.load(os.path.join(expert_traj_dir,'q.npy'))
        self.data_dq = np.load(os.path.join(expert_traj_dir,'dq.npy'))

        pos_list = []
        orn_list = []
        for i in range(len(self.data_q)):
          self.robot.setJointValue(self.data_q[i],gripper=220)
          pos = self.robot.getEndEffectorPos()
          pos_list.append(pos)
          orn = self.robot.getEndEffectorOrn()
          orn_list.append(orn)
        self.pos_traj = pos_list
        self.orn_traj = orn_list

        self.fix_orn = orn_list

        self.count = 0

    def init_obj(self):
        self.obj_file = os.path.join(self.urdf_dir,"objmodels/nut.urdf")
        self.obj_position = [0.35, 0.1, 0.31]
        self.obj_scaling = 2
        self.obj_orientation = self.p.getQuaternionFromEuler([math.pi/2, -math.pi/2, 0])
        self.obj_id = self.p.loadURDF(fileName=self.obj_file, basePosition=self.obj_position,baseOrientation=self.obj_orientation,globalScaling=self.obj_scaling)

    def reset_obj(self):
        obj_x = 0.35
        obj_y = 0.1
        transl = np.random.uniform(-0.1,0.1,size=(2,))
        self.obj_position = [obj_x + transl[0], obj_y + transl[1], 0.31]
        self.obj_orientation = self.p.getQuaternionFromEuler ([-math.pi / 2, 0, 0])
        self.p.resetBasePositionAndOrientation (self.obj_id, self.obj_position, self.obj_orientation)
        p.changeVisualShape (self.obj_id, -1, rgbaColor=[1, 0, 0, 1])
        for i in range(5):
          self.p.stepSimulation()

    def init_grasp(self):
        for j in range (7):
            self.p.resetJointState(self.robotId, j, self.data_q[0][j], self.data_dq[0][j])
        self.robot.gripperControl(0)

        self.null_q = self.data_q[0]
        orn = self.orn_traj[0]

        self.fix_orn = np.array(orn)
        self.fix_orn = np.expand_dims(orn,axis=0)

        pos = [self.obj_position[0]-0.03, self.obj_position[1]+0.2, self.obj_position[2] + 0.18]
        for i in range(30):
           self.robot.positionControl(pos,orn,null_pose=self.null_q,gripperPos=0)

        pos = [self.obj_position[0]-0.03, self.obj_position[1]+0.15, self.obj_position[2] + 0.075]
        for i in range(30):
           self.robot.positionControl(pos,orn,null_pose=self.null_q,gripperPos=0)

        pos = self.robot.getEndEffectorPos()
        pos[0] += np.random.uniform(-0.03,0.03)
        pos[1] += np.random.uniform(0,0.05)
        pos[2] += np.random.uniform(0,0.01)
        for i in range(30):
           self.robot.positionControl(pos,orn,null_pose=self.null_q,gripperPos=0)

        self.p.resetBasePositionAndOrientation (self.obj_id, self.obj_position, self.obj_orientation)

    def get_success(self,seg=None):
        pos = self.p.getBasePositionAndOrientation(self.obj_id)[0]
        if pos[1] < self.obj_position[1] - 0.1:
          self.count += 1
          r_pos = self.robot.getEndEffectorPos()
          return True
        else:
          return False


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import division, print_function

# Importing manupulation packages
from astropy.io import fits
import numpy as np
import glob
from numpy.polynomial import chebyshev
from scipy import ndimage
from scipy.signal import medfilt2d
from astropy.convolution import Gaussian1DKernel, Gaussian2DKernel, convolve

import sys
import argparse
import os

# Plotting
import matplotlib; matplotlib.use('TkAgg')
import matplotlib.pyplot as pl
import seaborn; seaborn.set_style('ticks')
import copy

from util import *

class XSHcomb:
    """
    Class to contain XSH spectrscopy combinations.
    """

    def __init__(self, list_of_files, base_name, sky, synth_sky, sky2d=None):
        """
        Instantiate fitsfiles. Input list of file-names to be combined.
        """

        if len(list_of_files) == 0:
            raise ValueError("Input file list empty")

        self.list_of_files = list_of_files
        self.list_of_skyfiles = sky

        fitsfile, header = {}, {}
        flux, error, bpmap = {}, {}, {}
        seeing = {}
        for ii, kk in enumerate(self.list_of_files):
            fitsfile[ii] = fits.open(kk)
            header[ii] = fitsfile[ii][0].header
            flux[ii] = fitsfile[ii][0].data
            error[ii] = fitsfile[ii][1].data
            bpmap[ii] = fitsfile[ii][2].data
            seeing[ii] = np.mean([header[ii]["HIERARCH ESO TEL AMBI FWHM START"], header[ii]["HIERARCH ESO TEL AMBI FWHM END"]])
            if sky2d is not None:
                flux[ii] += sky2d[ii]

        self.FWHM = np.median(list(seeing.values()))

        em_sky = []
        for ii, kk in enumerate(self.list_of_skyfiles):
            # fitsfile[ii] = fits.open(kk)
            em_sky.append(np.median(fits.open(kk)[0].data, axis = 0))

        self.fitsfile = fitsfile
        self.header = header
        self.flux = flux
        self.error = error
        self.bpmap = bpmap

        # Constructs WCS
        self.haxis = convert_air_to_vacuum(10.*(((np.arange(self.header[0]['NAXIS1'])) + 1 - self.header[0]['CRPIX1'])*self.header[0]['CDELT1']+self.header[0]['CRVAL1']))
        self.vaxis = (np.arange(self.header[0]['NAXIS2']) - self.header[0]['CRPIX2'])*self.header[0]['CDELT2']+self.header[0]['CRVAL2']

        if len(em_sky) == 0:
            print("No sky-frame given ... Using science image collapsed in the spatial direction ...")
            try:
                em_sky = np.sum(np.array(flux.values()), axis = 1)
            except:
                em_sky = None
        self.em_sky = em_sky
        self.base_name = base_name
        self.synth_sky = synth_sky

    def combine_imgs(self, NOD=False, same=False, repeats=1):
        """
        Combines X-shooter images.

        Function to inverse weighting combine multiple exposures from X-shooter. Tailored for combination of STARE-mode reduced images in NOD-sequence ABBA.

        Args:
            arg1 (bool): True to form nodding pairs before combination.

        Returns:
            fitsfile: fitsfile containing the combined flux, error and bad-pixel maps in consequtive extensions.

        """
        print("Combining "+str(len(self.list_of_files))+" files,\n"+str(self.list_of_files)+"\nto file:\n"+self.base_name+".fits....")
        img_nr = len(self.fitsfile)
        img_nr_list = np.arange(img_nr)

        pix_offsetx, pix_offsety = np.ones_like(img_nr_list), np.ones_like(img_nr_list)
        naxis1, naxis2 = np.ones_like(img_nr_list), np.ones_like(img_nr_list)
        exptimes = np.ones_like(img_nr_list)
        ra, dec = [0]*len(img_nr_list), [0]*len(img_nr_list)
        full_edge_mask = self.bpmap.copy()
        for ii, kk in enumerate(self.fitsfile):
            # Arrays to contain axis indices
            naxis1[ii] = self.header[ii]['NAXIS1']
            naxis2[ii] = self.header[ii]['NAXIS2']
            exptimes[ii] = self.header[ii]['EXPTIME']
            ra[ii], dec[ii] = float(self.header[ii]['RA']), float(self.header[ii]['DEC'])
        ref_ra, ref_dec = ra[0], dec[0]

        for ii, kk in enumerate(self.fitsfile):
            try:
                pix_offsetx[ii] = int(round(self.header[ii]['HIERARCH ESO SEQ CUMOFF X'] / self.header[ii]['CDELT1']))
                pix_offsety[ii] = int(round(self.header[ii]['HIERARCH ESO SEQ CUMOFF Y'] / self.header[ii]['CDELT2']))
            except KeyError:
                try:
                    # Offset mode along slit for older observations
                    from astropy.coordinates import SkyCoord, SkyOffsetFrame
                    import astropy.units as u
                    point_ra = self.header[ii]['RA']*u.deg
                    point_dec = self.header[ii]['DEC']*u.deg
                    offset_ra = self.header[ii]['HIERARCH ESO SEQ CUMOFF RA']*u.arcsec
                    offset_dec = self.header[ii]['HIERARCH ESO SEQ CUMOFF DEC']*u.arcsec
                    center = SkyCoord(ra=point_ra, dec=point_dec, frame = self.header[ii]['RADECSYS'].lower())
                    off_ra = point_ra + offset_ra
                    off_dec = point_dec + offset_dec
                    other = SkyCoord(ra=off_ra, dec=off_dec, frame = self.header[ii]['RADECSYS'].lower())
                    offset = center.separation(other).arcsecond
                    # Assume offset is along slit axis
                    pix_offsetx[ii] = int(round(0 / self.header[ii]['CDELT1']))
                    pix_offsety[ii] = int(round(offset / self.header[ii]['CDELT2']))

                except KeyError:
                    print("No header keyword: HIERARCH ESO SEQ CUMOFF X or HIERARCH ESO SEQ CUMOFF Y")
                    pix_offsetx[ii] = 0
                    pix_offsety[ii] = 0

            if same:

                # Wavelength step in velocity
                midwl = (max(self.haxis) - min(self.haxis))/2
                dv = 3e5*10*self.header[ii]['CDELT1']/midwl
                pix_offsetx[ii] = int(round((self.header[ii]['HIERARCH ESO QC VRAD BARYCOR']  + (self.header[ii]['WAVECORR']-1)*3e5)  / dv))
                # # Assume object is centered
                pix_offsety[ii] = int(round((max(naxis2)/2 - naxis2[ii]/2)))

            # Pixel numbers in x- and y-direction
            xs = np.arange(naxis1[ii]) + 1
            ys = np.arange(naxis2[ii]) + 1

            # Masking 1 pixel edges in frames.
            edge_len = 2
            if NOD:
                edge_len = 0
            edge_mask = (ys > max(ys) - edge_len) | (ys < min(ys) + edge_len)
            full_edge_mask[ii] = np.tile(edge_mask , (len(xs), 1)).T

        pix_offsetxmax = abs(max(pix_offsetx) - min(pix_offsetx))
        pix_offsetymax = abs(max(pix_offsety) - min(pix_offsety))
        # Defining size of out-array
        v_size = max(naxis1) + pix_offsetxmax
        h_size = max(naxis2) + pix_offsetymax

        # Data storage
        flux_cube = np.zeros((h_size, v_size, img_nr))
        error_cube = np.zeros((h_size, v_size, img_nr))
        bpmap_cube = np.ones((h_size, v_size, img_nr))

        # Manually mask bad region in VIS arm
        if self.header[ii]['HIERARCH ESO SEQ ARM'] == "VIS":
            for ii, kk in enumerate(img_nr_list):
                for xx, pp in enumerate(np.arange(11220, 11340, 1)):
                    self.bpmap[ii][int(round(26 - 0.2 * xx)):int(round(33 - 0.2 * xx)), pp] = 543

        for ii, kk in enumerate(img_nr_list):

            self.bpmap[ii] = self.bpmap[ii] + full_edge_mask[ii].astype("bool")*100

            # Defining positional offset between the frames.
            pos_v, pos_h = pix_offsety[kk], pix_offsetx[kk]  # offset

            # Finding the indices of the container in which to put image.
            offv = pix_offsety[kk] - min(pix_offsety)
            offh = pix_offsetx[kk] - min(pix_offsetx)

            # Define slices where to put image
            v_range1 = slice(offv, naxis2[ii] + offv)
            h_range1 = slice(offh, naxis1[ii] + offh)

            # b1 is full-size container with all values masked and b2 is input image with edge-mask + bad pixel mask.
            b1 = np.zeros((h_size, v_size))
            b2 = self.flux[ii]

            # Insert smaller (b3, input image) frame into larger frame (container)
            b1[v_range1, h_range1] = b2

            # Append to list containing flux images
            flux_cube[:, :, ii] = b1

            # Repeat for error extension
            b3 = np.zeros((h_size, v_size))
            b4 = self.error[ii]
            b3[v_range1, h_range1] = b4
            error_cube[:, :, ii] = b3

            # Repeat for bad pixel map
            b5 = np.ones((h_size, v_size))
            b6 = self.bpmap[ii]
            # Grow bap pixel regions !! Deprecated after update to pipeline version 2.8.3
            # b6 = np.rint(convolve(np.array(self.bpmap[ii]), Gaussian2DKernel(0.3)))
            b5[v_range1, h_range1] = b6
            bpmap_cube[:, :, ii] = b5

        # Mask 3-sigma outliers in the direction of the stack
        m, s = np.ma.median(np.ma.array(flux_cube, mask=bpmap_cube), axis = 2).data,  np.std(np.ma.array(flux_cube, mask=bpmap_cube), axis = 2).data
        if self.header[ii]['HIERARCH ESO SEQ ARM'] == "NIR":
            sigma_mask = 5
        else:
            sigma_mask = 3
        l, h = np.tile((m - sigma_mask*s).T, (img_nr, 1, 1)).T, np.tile((m + sigma_mask*s).T, (img_nr, 1, 1)).T
        bpmap_cube[(flux_cube < l) | (flux_cube > h)] = 666

        # Form nodding pairs
        if NOD:
            if not repeats == 1:
                # Smaller container
                flux_cube_tmp = np.zeros((h_size, v_size, int(np.ceil(img_nr / repeats))))
                error_cube_tmp = np.zeros((h_size, v_size, int(np.ceil(img_nr / repeats))))
                bpmap_cube_tmp = np.zeros((h_size, v_size, int(np.ceil(img_nr / repeats))))
                # Collapse in repeats
                for ii, kk in enumerate(np.arange(int(np.ceil(img_nr / repeats)))):
                    # Make lower an upper index of files, which is averaged over. If all NOD positions does not have the same number of repeats, assume the last position is cut.
                    low, up = ii*repeats, min(img_nr, (ii+1)*repeats)
                    # Slice structure
                    subset = slice(low, up)
                    # Average over subset
                    flux_cube_tmp[:, :, ii], error_cube_tmp[:, :, ii], bpmap_cube_tmp[:, :, ii] = avg(flux_cube[:, :, subset], error_cube[:, :, subset], bpmap_cube[:, :, subset].astype("bool"), axis=2)
                # Update number holders
                img_nr_list = np.arange(img_nr/repeats)
                pix_offsety = pix_offsety[::repeats]
                flux_cube, error_cube, bpmap_cube = flux_cube_tmp, error_cube_tmp, bpmap_cube_tmp

            # Form the pairs [(A1-B1) - shifted(B1-A1)] and [(B2-A2) - shifted(A2-B2)] at positions 0, 2. Sets the other images to np.nan.
            flux_cube, error_cube, bpmap_cube = form_nodding_pairs(flux_cube, error_cube,  bpmap_cube, max(naxis2), pix_offsety)

        # Introduce filter based on smoothed image
        flux_filt, error_filt, bp_filt = np.zeros_like(error_cube), np.zeros_like(error_cube), np.zeros_like(error_cube)
        for ii in range(error_cube.shape[2]):
            flux_filt[:, :, ii] = medfilt2d(flux_cube[:, :, ii], 3)
            error_filt[:, :, ii] = medfilt2d(error_cube[:, :, ii], 3)

        bp_filt = abs(flux_filt - flux_cube) > 7 * error_filt
        bp_filt[int(np.floor(h_size/2-2)):int(np.ceil(h_size/2+2)), :, :] = False
        bpmap_cube[bp_filt == True] = 555

        # Mask outliers

        bpmap_cube[(flux_cube == 0) | (flux_cube == 1) | (flux_cube == np.inf) | (flux_cube == np.nan)] = 666
        bpmap_cube[(error_cube == 0) | (error_cube == 1) | (error_cube == np.inf) | (error_cube == np.nan)] = 666

        # Boolean mask based on the bad-pixel map, the edge mask and the sigma-clipped mask
        mask_cube = (bpmap_cube != 0)

        # Make weight map based on background variance in boxcar window
        shp = error_cube.shape
        weight_cube = np.zeros_like(error_cube)
        if not same:
            for ii in range(shp[2]):
                run_var = np.ones(shp[1])
                for kk in np.arange(shp[1]):
                    err_bin = 1/(error_cube[:, kk-4:kk+4, ii][~(bpmap_cube[:, kk-4:kk+4, ii].astype("bool"))])**2
                    if len(err_bin) != 0:
                        run_var[kk] = np.median(err_bin.flatten())
                weight_cube[:, :, ii] = np.tile(run_var, (shp[0], 1))
        elif same:
            for ii in range(shp[2]):
                weight_cube[:, :, ii] = np.tile(np.median(1/(error_cube[~(bpmap_cube.astype("bool"))])**2), (shp[0], shp[1]))

        # Normlize weights
        weight_cube[mask_cube] = 0
        weight_cube = weight_cube/np.tile(np.sum(weight_cube, axis=2).T, (shp[2], 1, 1)).T

        # Calculate mean and error
        mean, error, bpmap = avg(flux_cube, error_cube, mask_cube, axis=2, weight_map = weight_cube)

        # Assign new flux and error
        mean[np.isnan(mean)] = 0
        self.flux = mean
        self.error = error
        self.bpmap = bpmap

        if same:
            self.flux[np.isnan(self.flux)] = np.median(self.flux[~np.isnan(self.flux)])
            self.error[np.isnan(self.error)] = 10. * max(self.error[~np.isnan(self.error)])

        # Write to file
        wrf = np.where(pix_offsety == min(pix_offsety))[0][0]

        self.fitsfile = self.fitsfile[wrf]
        self.header = self.header[wrf]
        self.fitsfile[0].data = self.flux.data
        self.fitsfile[1].data = self.error.data
        self.fitsfile[2].data = self.bpmap.data

        # Update WCS
        self.header["CRVAL2"] = self.header["CRVAL2"] - (max(pix_offsety - min(pix_offsety)))  * self.header["CDELT2"]

        # Set header keyword based on median seeing
        self.header["SEEING"] = self.FWHM

        if not same:
            # Simply combined image
            if not NOD:
                # Updating file header
                self.fitsfile.header = self.header
                # Updating extention header keywords
                self.fitsfile[1].header["CRVAL2"], self.fitsfile[2].header["CRVAL2"] = self.fitsfile[0].header["CRVAL2"], self.fitsfile[0].header["CRVAL2"]
                self.fitsfile.writeto(self.base_name+".fits", overwrite =True)
            # If nodded
            elif NOD:
                # self.header["WAVECORR"] = self.correction_factor
                self.fitsfile.header = self.header
                self.fitsfile[1].header["CRVAL2"], self.fitsfile[2].header["CRVAL2"] = self.fitsfile[0].header["CRVAL2"], self.fitsfile[0].header["CRVAL2"]
                self.fitsfile.writeto(self.base_name+"skysub.fits", overwrite =True)
        elif same:
            self.fitsfile.header = self.header
            self.fitsfile[1].header["CRVAL2"], self.fitsfile[2].header["CRVAL2"] = self.fitsfile[0].header["CRVAL2"], self.fitsfile[0].header["CRVAL2"]
            self.base_name = self.base_name[:-3]+"_combined"
            self.fitsfile.writeto(self.base_name+".fits", overwrite =True)

        # Update WCS axis
        self.vaxis = (np.arange(self.header['NAXIS2']) - self.header['CRPIX2'])*self.header['CDELT2']+self.header['CRVAL2']

        mask = (self.flux > -1e-17) & (self.flux < 1e-17)
        hs, ed = np.histogram(self.flux[mask], bins=1000000)
        print("")
        print("Mode of flux values (should be close to zero):")
        print(ed[find_nearest(hs, max(hs))], ed[find_nearest(hs, max(hs))+ 1])
        print("")
        print("Finished combining files ...")

        # Get binned spectrum
        bin_length = int(len(self.haxis) / 300)
        bin_flux, bin_error, __ = bin_image(self.flux, self.error, self.bpmap, bin_length, weight = True)

        # Save binned image for quality control
        self.fitsfile[0].data = bin_flux
        self.fitsfile[1].data = bin_error
        self.fitsfile[0].header["CD1_1"] = self.fitsfile[0].header["CD1_1"] * bin_length
        self.fitsfile.writeto(self.base_name+"_binned.fits", overwrite=True)
        self.fitsfile[0].data = self.flux
        self.fitsfile[1].data = self.error
        self.fitsfile[0].header["CD1_1"] = self.fitsfile[0].header["CD1_1"] / bin_length

    def sky_subtract(self, seeing, masks, sky_check=False, nod=False):

        """Sky-subtracts X-shooter images.

        Function to subtract sky off a rectifed x-shooter image. Fits a low-order polynomial in the spatial direction at each pixel and subtracts this off the same pixel. Assumes the object is centrered and masks the trace according to the seeing. Additionally accepts a list of arcsecond offsets with potential additional objects in the slit.

        fitsfile : fitsfile
            Input image to be sky-subtracted
        seeing : float
            Seeing of observation used to mask trace. Uses this as the width from the center which is masked. Must be wide enough to completely mask any signal.
        masks : list
            List of floats to additionally mask. This list contains the offset in arcsec relative to the center of the slit in which additional sources appear in the slit and which must be masked for sky-subtraction purposes. Positive values is defined relative to the WCS.

        Returns
        -------
        sky-subtracted image : fitsfile

        Notes
        -----
        na
        """
        if not nod:
            print("")
            print('Subtracting sky....')
            # Make trace mask
            seeing_pix = seeing / self.header['CDELT2']
            trace_offsets = np.array(masks) / self.header['CDELT2']
            traces = []
            for ii in trace_offsets:
                traces.append(self.header['NAXIS2']/2 + ii)

            # Masking pixels in frame.
            trace_mask = np.zeros(self.header['NAXIS2']).astype("bool")
            for ii, kk in enumerate(traces):
                trace_mask[int(kk - seeing_pix):int(kk + seeing_pix)] = True
            full_trace_mask = np.tile(trace_mask , (self.header['NAXIS1'], 1)).T
            full_mask = self.bpmap.astype("bool") | full_trace_mask

            sky_background = np.zeros_like(self.flux)
            sky_background_error = np.zeros_like(self.flux)
            for ii, kk in enumerate(self.haxis):
                # Pick mask slice
                mask = full_mask[:, ii]

                # Sigma clip before sky-estimate to remove noisy pixels with bad error estimate.
                m, s = np.nanmean(self.flux[:, ii]), np.nanstd(self.flux[:, ii])
                clip_mask = (self.flux[:, ii] < m - s) | (self.flux[:, ii] > m + s)

                # Combine masks
                mask = mask | clip_mask

                # Subtract polynomial estiamte of sky
                vals = self.flux[:, ii][~mask]
                errs = self.error[:, ii][~mask]

                try:
                    chebfit = chebyshev.chebfit(self.vaxis[~mask], vals, deg = 2, w=1/errs)
                    chebfitval = chebyshev.chebval(self.vaxis, chebfit)
                    chebfitval[chebfitval <= 0] = 0
                except TypeError:
                    print("Empty array for sky-estimate at index "+str(ii)+". Sky estimate replaced with zeroes.")
                    chebfitval = np.zeros_like(self.vaxis)
                except:
                    print("Polynomial fit did not converge at index "+str(ii)+". Sky estimate replaced with median value.")
                    chebfitval = np.ones_like(self.vaxis)*np.ma.median(self.vaxis[~mask])

                if ii % int(self.header['NAXIS1']/5) == 0 and sky_check and ii != 0:
                    # Plotting for quality control
                    pl.errorbar(self.vaxis[~mask], vals, yerr=errs, fmt=".k", capsize=0, elinewidth=0.5, ms=3)
                    pl.plot(self.vaxis, chebfitval)
                    pl.xlabel("Spatial index")
                    pl.ylabel("Flux density")
                    pl.title("Quality test: Sky estimate at index: "+str(ii) )
                    pl.show()

                sky_background[:, ii] = chebfitval

            # Subtract sky
            self.flux = self.flux - convolve(sky_background, Gaussian2DKernel(1.0))

        self.em_sky = np.sum(self.em_sky, axis=0)
        # Calibrate wavlength solution
        XSHcomb.finetune_wavlength_solution(self)
        self.sky_mask = np.tile(self.sky_mask, (self.header["NAXIS2"], 1)).astype("int")

        self.bpmap += self.sky_mask
        self.flux[self.bpmap.astype("bool")] = 0
        self.fitsfile.header = self.header
        self.fitsfile[0].data, self.fitsfile[1].data = self.flux.data, self.error.data
        self.fitsfile[2].data = self.bpmap.data
        self.fitsfile.writeto(self.base_name+"skysub.fits", overwrite =True)

        print('Writing sky subtracted image to '+self.base_name+"skysub.fits")

    def finetune_wavlength_solution(self):
        print("")
        print("Cross correlating with synthetic sky to obtain refinement to wavlength solution ...")
        print("")

        # Remove continuum
        mask = ~np.isnan(self.em_sky) & (self.haxis < 18000) & (self.haxis > 3500)
        hist, edges = np.histogram(self.em_sky[mask], bins="auto")
        max_idx = find_nearest(hist, max(hist))
        sky = self.em_sky - edges[max_idx]
        mask = ~np.isnan(sky) & (sky > 0) & (self.haxis < 18000) & (self.haxis > 3500)

        # Load synthetic sky
        sky_model = fits.open("statics/skytable_hres.fits")
        wl_sky = 1e4*(sky_model[1].data.field('lam')) # In micron
        flux_sky = sky_model[1].data.field('flux')

        # Convolve to observed grid
        from scipy.interpolate import interp1d
        f = interp1d(wl_sky, convolve(sky_model[1].data.field('flux'), Gaussian1DKernel(stddev=10)), bounds_error=False, fill_value=np.nan)
        synth_sky = f(self.haxis)

        # Cross correlate with redshifted spectrum and find velocity offset
        offsets = np.arange(-0.0005, 0.0005, 0.00001)
        correlation = np.zeros(offsets.shape)
        for ii, kk in enumerate(offsets):
            synth_sky = f(self.haxis * (1. + kk))
            correlation[ii] = np.correlate(sky[mask]*(np.nanmax(synth_sky)/np.nanmax(sky)), synth_sky[mask])

        # Index with maximal value
        max_idx = find_nearest(correlation, max(correlation))
        # Corrections to apply to original spectrum, which maximizes correlation.
        self.correction_factor = 1. + offsets[max_idx]
        print("Found preliminary velocity offset: "+str((self.correction_factor - 1.)*3e5)+" km/s")
        print("")
        print("Minimising residuals between observed sky and convolved synthetic sky to obtain the sky PSF ...")
        print("")

        # Zero-deviation wavelength of arms, from http://www.eso.org/sci/facilities/paranal/instruments/xshooter/doc/VLT-MAN-ESO-14650-4942_v87.pdf
        if self.header['HIERARCH ESO SEQ ARM'] == "UVB":
            zdwl = 4050
            pixel_width = 50
        elif self.header['HIERARCH ESO SEQ ARM'] == "VIS":
            zdwl = 6330
            pixel_width = 50
        elif self.header['HIERARCH ESO SEQ ARM'] == "NIR":
            zdwl = 13100
            pixel_width = 50

        # Get seeing PSF by minimizing the residuals with a theoretical sky model, convovled with an increasing psf.
        psf_width = np.arange(1, pixel_width, 1)
        res = np.zeros(psf_width.shape)
        for ii, kk in enumerate(psf_width):
            # Convolve sythetic sky with Gaussian psf
            convolution = convolve(flux_sky, Gaussian1DKernel(stddev=kk))
            # Interpolate high-res syntheric sky onto observed wavelength grid.
            f = interp1d(wl_sky, convolution, bounds_error=False, fill_value=np.nan)
            synth_sky = f(self.haxis*self.correction_factor)
            # Calculate squared residuals
            residual = np.nansum((synth_sky[mask]*(np.nanmax(sky[mask])/np.nanmax(synth_sky[mask])) - sky[mask])**2.)
            res[ii] = residual

        # Index of minimal residual
        min_idx = find_nearest(res, min(res))

        # Wavelegth step corresponding psf width in FWHM
        R, seeing = np.zeros(psf_width.shape), np.zeros(psf_width.shape)
        for ii, kk in enumerate(psf_width):
            dlambda = np.diff(wl_sky[::kk])*2*np.sqrt(2*np.log(2))
            # Interpolate to wavelegth grid
            f = interp1d(wl_sky[::kk][:-1], dlambda, bounds_error=False, fill_value=np.nan)
            dlambda = f(self.haxis)

            # PSF FWHM in pixels
            d_pix = dlambda/(10*self.header['CDELT1'])
            # Corresponding seeing PSF FWHM in arcsec
            spatial_psf = d_pix*self.header['CDELT2']

            # Index of zero-deviation
            zd_idx = find_nearest(self.haxis, zdwl)

            # Resolution at zero-deviation wavelength
            R[ii] = (self.haxis/dlambda)[zd_idx]
            # Seeing at zero-deviation wavelength
            seeing[ii] = spatial_psf[zd_idx]

        fig, ax1 = pl.subplots()
        ax1.yaxis.set_major_formatter(pl.NullFormatter())

        convolution = convolve(flux_sky, Gaussian1DKernel(stddev=psf_width[min_idx]))
        f = interp1d(wl_sky, convolution, bounds_error=False, fill_value=np.nan)
        synth_sky = f(self.haxis)

        # Cross correlate with redshifted spectrum and find velocity offset
        offsets = np.arange(-0.0005, 0.0005, 0.000001)
        correlation = np.zeros_like(offsets)
        for ii, kk in enumerate(offsets):
            synth_sky = f(self.haxis * (1. + kk))
            correlation[ii] = np.correlate(sky[mask]*(np.nanmax(synth_sky[mask])/np.nanmax(sky[mask])), synth_sky[mask])

        # Smooth cross-correlation
        correlation = convolve(correlation, Gaussian1DKernel(stddev=20))

        # Index with maximum correlation
        max_idx = find_nearest(correlation, max(correlation))
        self.correction_factor = (1. + offsets[max_idx])
        self.header["WAVECORR"] = self.correction_factor
        print("Found refined velocity offset: "+str((self.correction_factor - 1.)*3e5)+" km/s")
        print("")

        # Mask flux with > 3-sigma sky brightness
        self.sky_mask = f(self.haxis*self.correction_factor) > np.percentile(f(self.haxis*self.correction_factor), 99)
        ax2 = ax1.twiny()

        ax2.errorbar(offsets[max_idx]*3e5, max(correlation)*(max(res)/max(correlation)), fmt=".k", capsize=0, elinewidth=0.5, ms=13, label="Wavelength correction:" + str(np.around((self.correction_factor - 1.)*3e5, decimals = 1)) +" km/s", color="r")
        ax2.plot(offsets*3e5, correlation*(max(res)/max(correlation)), color="r")
        ax2.set_xlabel("Offset velocity / [km/s]", color="r")
        ax2.set_ylabel("Cross correlation", color="r")
        ax2.yaxis.set_label_position("right")
        ax2.yaxis.set_major_formatter(pl.NullFormatter())
        ax2.legend(loc=2)
        pl.savefig(self.base_name+"Wavelength_cal.pdf")
        pl.clf()

def run_combination(args):
    # Load in files
    sky2d = None
    response_2d = 1
    if args.mode == "STARE" or args.mode == "NOD":
        files = sorted(glob.glob(args.filepath+"reduced_data/"+args.OB+"/"+args.arm+"/*/*SCI_SLIT_MERGE2D_*.fits"))
        if len(files) == 0:
            print("No files found... Exitting..")
            exit()
        sky_files = sorted(glob.glob(args.filepath+"reduced_data/"+args.OB+"/"+args.arm+"/*/*SKY_SLIT_MERGE2D_*.fits"))
        n_flux_files = len(sorted(glob.glob(args.filepath+"reduced_data/"+args.OB+"/"+args.arm+"/*/*SCI_SLIT_FLUX_MERGE2D_*.fits")))
        if n_flux_files == 0 and not args.use_master_response:
            print("Option to use master response function has not been set and no flux calibrated data exists. You should probably set the optional argument, --use_master_response.")
            print("Press \"enter\" to continue anyway and use the non-flux calibrated images.")
            raw_input()
        if not args.use_master_response and n_flux_files != 0:
            response_2d = [fits.open(ii)[0].data for ii in files]

            merge_files = glob.glob(args.filepath+"reduced_data/"+args.OB+"/"+args.arm+"/*/*/*FLUX_ORDER2D*")
            target_files = glob.glob(args.filepath+"reduced_data/"+args.OB+"/"+args.arm+"/*/*FLUX_MERGE2D*")
            target_files = [ii for ii in target_files if "MANMERGE" not in ii and "TELL" not in ii]
            for kk, ll in list(zip(merge_files, target_files)):
                insorder2d = XshOrder2D(kk)
                fname_out = ll.replace('FLUX_MERGE2D', 'FLUX_MERGE2D_MANMERGE')
                insorder2d.do_all(fname_out)

            files = sorted(glob.glob(args.filepath+"reduced_data/"+args.OB+"/"+args.arm+"/*/*MANMERGE_*.fits"))

            if len(files) == 0:
                files = sorted(glob.glob(args.filepath+"reduced_data/"+args.OB+"/"+args.arm+"/*/*SCI_SLIT_FLUX_MERGE2D_"+args.arm+".fits"))

            response_2d = [np.tile(np.nanmedian(fits.open(kk)[0].data/response_2d[ii], axis=0), (np.shape(response_2d[ii])[0], 1)) for ii, kk in enumerate(files)]

            np.savetxt(args.filepath+"reduced_data/"+args.OB+"/"+args.arm+"/response_function.dat", np.nanmean(np.nanmean(response_2d, axis=1), axis=0))

        if args.mode == "NOD":
            sky2d = sorted(glob.glob(args.filepath+"reduced_data/"+args.OB+"/"+args.arm+"/*/*SKY_SLIT_MERGE2D_*.fits"))
            sky2d = np.array([fits.open(ii)[0].data for ii in sky2d]) * np.array(response_2d)
    elif args.mode == "COMBINE":

        files = sorted(glob.glob(args.filepath+args.arm+"*skysub.fits"))
        if len(files) == 0:
            print("No files found... Exitting..")
            exit()
        sky_files = sorted(glob.glob(args.filepath+"reduced_data/"+args.OB+"/"+args.arm+"/*/*SKY_SLIT_MERGE2D_*.fits"))

    skyfile = sorted(glob.glob("statics/"+args.arm+"skytable.fits"))

    img = XSHcomb(files, args.filepath+args.arm+args.OB, sky=sky_files, synth_sky=skyfile, sky2d=sky2d)
    # Combine nodding observed pairs.
    if args.mode == "STARE":
        img.combine_imgs(NOD=False)
        img.sky_subtract(seeing=args.seeing, masks=args.masks, sky_check=False)
    elif args.mode == "NOD":
        img.combine_imgs(NOD=True, repeats=args.repeats)
        img.sky_subtract(seeing=args.seeing, masks=args.masks, sky_check=False, nod=True)
    elif args.mode == "COMBINE":
        img.combine_imgs(same=True)

def main(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('filepath', type=str, default="/Users/jonatanselsing/Work/work_rawDATA/XSGRB/GRB120327A/", help='Path to burst directory on which to run combination. Directory must contain a centain directory structure, similar to /Users/jonatanselsing/Work/work_rawDATA/XSGRB/GRB121027A/reduced_data/OB1/UVB/XSHOO.2012-10-30T05:03:26.098cosmicced/ToO_GRBtrigger_4x600_SCI_SLIT_FLUX_MERGE2D_UVB.fits.')
    parser.add_argument('arm', type=str, default="UVB", help='X-shooter arm to combine. Used to find files')
    parser.add_argument('mode', type=str, default="STARE", help='MODE in which to run combinations. Can either be STARE, NOD or COMBINE')
    parser.add_argument('OB', type=str, default="OB1", help='OB number. Used to look for files.')
    parser.add_argument('-repeats', type=int, default=1, help='Number of times nodding position has been repeated')
    parser.add_argument('-seeing', type=float, default=1.0, help='Estimated seeing. Used to mask trace for sky-subtraction.')
    parser.add_argument('-masks', type=list, default=list(), help='List of offsets relative to center of additional masks for sky-subtraction.')
    parser.add_argument('--use_master_response', action="store_true" , help = 'Set this optional keyword if input files are not flux-calibrated. Used in sky-subtraction.')

    args = parser.parse_args(argv)

    if not args.filepath:
        print('When using arguments, you need to supply a filepath. Stopping execution')
        exit()

    print("Running combination on files: " + args.filepath)

    run_combination(args)

if __name__ == '__main__':
    # If script is run from editor or without arguments, run using this:
    if len(sys.argv) == 1:
        """
        Central scipt to combine images from X-shooter for the X-shooter GRB sample.
        """

        # Load in files
        parser = argparse.ArgumentParser()
        args = parser.parse_args()

        object_name = "/Users/jonatanselsing/Work/work_rawDATA/Bertinoro/"

        args.filepath = object_name

        arms = ["UVB"] # UVB, VIS, NIR, ["UVB", "VIS", "NIR"]

        combine = False # True False

        OBs = ["OB1"] # ["OB1", "OB2", "OB3", "OB4", "OB5", "OB6", "OB7", "OB8", "OB9", "OB10", "OB11", "OB12", "OB13", "OB14"]
        for ll in OBs:
            args.OB = ll
            # print(ll)
            for ii in arms:
                args.arm = ii # UVB, VIS, NIR
                # args.mode = "STARE"
                args.mode = "NOD"
                # if ii == "NIR":
                #     args.mode = "NOD"
                # if combine:
                #     args.mode = "COMBINE"

                args.use_master_response = False # True False
                args.masks = []
                args.seeing = 1.5
                args.repeats = 1

                run_combination(args)

    else:
        main(argv = sys.argv[1:])

reason: import numpy,from numpy,from scipy,import astropy,from astropy


text:
# This function concatenate a bunch of files into a large, single file for dataset reading

import os
import numpy as np
import pandas as pd

mother_folder = '/home/sr365/MM_Bench/Data/Omar_bowtie'
geometry, spectra = None, None

for file_prefix in ['inputs_Srico_MMPA_bowtie_x1000_','Abs_Srico_MMPA_bowtie_x1000_']:
    value_all = None        # Init the value
    for i in range(1, 8):
        # Combine the values
        file = os.path.join(mother_folder, file_prefix + '{}.csv'.format(i))
        value = pd.read_csv(file, sep=',', header=None).values
        #print(np.shape(value))
        if value_all is None:
            value_all = value
        else:
            value_all = np.concatenate([value_all, value])      # Concate the value
        print(np.shape(value_all))

    # Save the values
    np.savetxt(os.path.join(mother_folder, file_prefix + 'all.csv'), value_all)


text:
#!/usr/bin/env python
# coding: utf-8

import numpy as np
import sys
import umap
import warnings
from scipy import sparse
import networkx as nx
warnings.filterwarnings('ignore')
import seaborn as sns;sns.set()
import matplotlib.pyplot as plt
from sklearn.preprocessing import normalize
import pandas as pd
############################################
W = sparse.load_npz(sys.argv[1]) # adj.npz
txtfilename = sys.argv[2] # 'data.txt.gz'

degree = W.sum(axis=1) #calculate degree vector
WW = W.tocoo(copy=True)
norma = W.sum()
rowdegree = np.asarray([degree[ind] for ind in WW.row]).squeeze()
coldegree = np.asarray([degree[ind] for ind in WW.col]).squeeze()
datamodel = rowdegree*coldegree*1.0/norma
nullmodel = sparse.csr_matrix((datamodel, (WW.row, WW.col)), shape=W.shape)
M = W - nullmodel
modularity = M.sum(axis=1)

np.savetxt(txtfilename+'.nn'+str(10)+'.modularity.gz', modularity)

max_stars_repo_name: BiCroLab/WSI-analysis

max_issues_repo_name: BiCroLab/WSI-analysis

max_forks_repo_name: BiCroLab/WSI-analysis

reason: import numpy,from scipy,import networkx


text:
import copy
import datetime
import tempfile
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pytest

from imucal import FerrarisCalibrationInfo, TurntableCalibrationInfo
from imucal.management import CalibrationWarning, save_calibration_info
from nilspodlib.calibration_utils import (
    save_calibration,
    find_calibrations_for_sensor,
    find_closest_calibration_to_date,
)

# TODO: Try to remove tests that are already in imucal

@pytest.fixture()
def dummy_cal_dict():
    sample_data = {
        "K_a": np.identity(3),
        "R_a": np.identity(3),
        "b_a": np.zeros(3),
        "K_g": np.identity(3),
        "R_g": np.identity(3),
        "K_ga": np.identity(3),
        "b_g": np.zeros(3),
    }
    return sample_data

@pytest.fixture()
def dummy_cal(dummy_cal_dict):
    return FerrarisCalibrationInfo(**dummy_cal_dict)

def test_save_calibration(dummy_cal):
    with tempfile.TemporaryDirectory() as f:
        save_calibration(dummy_cal, "test", datetime.datetime(2000, 10, 3, 13, 22), f)

        assert next(Path(f).glob("*")).name == "test_2000-10-03_13-22.json"

@pytest.mark.parametrize("sensor_id", ["12345", "tes*", "b da", "tes"])
def test_save_cal_id_validation(dummy_cal, sensor_id):
    with tempfile.TemporaryDirectory() as f:
        with pytest.raises(ValueError):
            save_calibration(dummy_cal, sensor_id, datetime.datetime(2000, 10, 3, 13, 22), f)

@dataclass(eq=False)
class CustomFerraris(FerrarisCalibrationInfo):
    CAL_TYPE = "Custom Ferraris"
    custom_field: str = "default_custom_value"

@pytest.fixture()
def dummy_cal_folder(dummy_cal):
    with tempfile.TemporaryDirectory() as f:
        for sid in ["tes1", "tes2", "tes3"]:
            for min in range(10, 30, 2):
                date = datetime.datetime(2000, 10, 3, 13, min)
                save_calibration_info(dummy_cal, sid, date, f, folder_structure="")
        yield f

@pytest.fixture()
def dummy_cal_folder_recursive(dummy_cal_dict):
    with tempfile.TemporaryDirectory() as f:
        for s in [FerrarisCalibrationInfo, TurntableCalibrationInfo, CustomFerraris]:
            for sid in ["tes1", "tes2", "tes3"]:
                new_cal = s(**dummy_cal_dict)
                for min in range(10, 30, 2):
                    date = datetime.datetime(2000, 10, 3, 13, min)
                    save_calibration_info(new_cal, sid, date, Path(f))
        yield f

def test_find_calibration_simple(dummy_cal_folder):
    cals = find_calibrations_for_sensor("tes1", dummy_cal_folder)

    assert len(cals) == 10
    assert all(["tes1" in str(x) for x in cals])

def test_find_calibration_non_existent(dummy_cal_folder):
    with pytest.raises(ValueError):
        find_calibrations_for_sensor("wrong_sensor", dummy_cal_folder)

    cals = find_calibrations_for_sensor("wrong_sensor", dummy_cal_folder, ignore_file_not_found=True)

    assert len(cals) == 0

def test_find_calibration_recursive(dummy_cal_folder_recursive):
    with pytest.raises(ValueError):
        find_calibrations_for_sensor("tes1", dummy_cal_folder_recursive, recursive=False)

    cals = find_calibrations_for_sensor("tes1", dummy_cal_folder_recursive, recursive=True)

    assert len(cals) == 30
    assert all(["tes1" in str(x) for x in cals])

def test_find_calibration_type_filter(dummy_cal_folder_recursive):
    cals = find_calibrations_for_sensor("tes1", dummy_cal_folder_recursive, recursive=True, filter_cal_type="ferraris")

    assert len(cals) == 10
    assert all(["tes1" in str(x) for x in cals])

def test_find_calibration_type_filter_case_sensitive(dummy_cal_folder_recursive):
    cals = find_calibrations_for_sensor("tes1", dummy_cal_folder_recursive, recursive=True, filter_cal_type="ferraris")

    assert len(cals) == 10
    assert all(["tes1" in str(x) for x in cals])

    cals = find_calibrations_for_sensor("tes1", dummy_cal_folder_recursive, recursive=True, filter_cal_type="Ferraris")

    assert len(cals) == 10
    assert all(["tes1" in str(x) for x in cals])

    cals = find_calibrations_for_sensor("tes1", dummy_cal_folder_recursive, recursive=True, filter_cal_type="FERRARIS")

    assert len(cals) == 10
    assert all(["tes1" in str(x) for x in cals])

def test_find_closest(dummy_cal_folder):
    cal = find_closest_calibration_to_date("tes1", datetime.datetime(2000, 10, 3, 13, 14), dummy_cal_folder)

    assert cal.name == "tes1_2000-10-03_13-14.json"

    # Test that before and after still return the correct one if there is an exact match
    cal = find_closest_calibration_to_date(
        "tes1", datetime.datetime(2000, 10, 3, 13, 14), dummy_cal_folder, before_after="before"
    )

    assert cal.name == "tes1_2000-10-03_13-14.json"

    cal = find_closest_calibration_to_date(
        "tes1", datetime.datetime(2000, 10, 3, 13, 14), dummy_cal_folder, before_after="after"
    )

    assert cal.name == "tes1_2000-10-03_13-14.json"

def test_find_closest_non_existend(dummy_cal_folder):
    with pytest.raises(ValueError):
        find_closest_calibration_to_date("wrong_sensor", datetime.datetime(2000, 10, 3, 13, 14), dummy_cal_folder)

    cal = find_closest_calibration_to_date(
        "wrong_sensor", datetime.datetime(2000, 10, 3, 13, 14), dummy_cal_folder, ignore_file_not_found=True
    )

    assert cal is None

def test_find_closest_before_after(dummy_cal_folder):
    # Default to earlier if same distance before and after.
    cal = find_closest_calibration_to_date("tes1", datetime.datetime(2000, 10, 3, 13, 15), dummy_cal_folder)

    assert cal.name == "tes1_2000-10-03_13-14.json"

    # Return later if after.
    cal = find_closest_calibration_to_date(
        "tes1", datetime.datetime(2000, 10, 3, 13, 15), dummy_cal_folder, before_after="after"
    )

    assert cal.name == "tes1_2000-10-03_13-16.json"

    # Return later if before.
    cal = find_closest_calibration_to_date(
        "tes1", datetime.datetime(2000, 10, 3, 13, 15), dummy_cal_folder, before_after="before"
    )

    assert cal.name == "tes1_2000-10-03_13-14.json"

def test_find_closest_warning(dummy_cal_folder):
    with pytest.warns(CalibrationWarning) as rec:
        find_closest_calibration_to_date(
            "tes1", datetime.datetime(2000, 10, 3, 13, 15), dummy_cal_folder, warn_thres=datetime.timedelta(seconds=30)
        )

    assert len(rec) == 1

    with pytest.warns(None) as rec:
        find_closest_calibration_to_date(
            "tes1", datetime.datetime(2000, 10, 3, 13, 14), dummy_cal_folder, warn_thres=datetime.timedelta(seconds=30)
        )

    assert len(rec) == 0

    # test default
    with pytest.warns(None) as rec:
        find_closest_calibration_to_date("tes1", datetime.datetime(2000, 10, 3, 13, 15), dummy_cal_folder)

    assert len(rec) == 0

# Comment in once ref cals are open source
# def test_find_default_cal():
#     cals = find_calibrations_for_sensor("3d73")
#
#     assert len(cals) > 0
#
#
# def test_find_default_cal_wrong():
#     with pytest.raises(ValueError):
#         find_calibrations_for_sensor("FFFF")

max_stars_repo_name: mad-lab-fau/NilsPodPythonLib

max_issues_repo_name: mad-lab-fau/NilsPodPythonLib

max_forks_repo_name: mad-lab-fau/NilsPodPythonLib