text:
import os
import uuid
import time
import arrow
from skyportal.tests import api
from glob import glob
import yaml
import numpy as np
import datetime

from skyportal.enum_types import ALLOWED_SPECTRUM_TYPES, default_spectrum_type

def test_spectrum_put(super_admin_user, super_admin_token, public_source, lris):
    # make groups that must be unique to this test
    status, data = api(
        "POST",
        "groups",
        data={
            "name": str(uuid.uuid4()),
            "group_admins": [super_admin_user.id],
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data["status"] == "success"
    group_id1 = data["data"]["id"]

    status, data = api(
        "POST",
        "groups",
        data={
            "name": str(uuid.uuid4()),
            "group_admins": [super_admin_user.id],
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data["status"] == "success"
    group_id2 = data["data"]["id"]

    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': public_source.id,
            'observed_at': '2020-01-10T00:00:00',
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.3, 232.1, 235.3],
            'group_ids': [group_id1],
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    spectrum_id = data['data']['id']

    assert status == 200
    assert data['status'] == 'success'

    # update only the label
    custom_label = str(uuid.uuid4())
    status, data = api(
        'PUT',
        f'spectrum/{spectrum_id}',
        data={
            'label': custom_label,
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    status, data = api(
        'GET',
        f'spectrum/{spectrum_id}',
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['label'] == custom_label
    group_ids = [g['id'] for g in data['data']['groups']]
    assert group_id1 in group_ids
    assert group_id2 not in group_ids

    # update the group IDs (should ADD group2, not remove group1)
    status, data = api(
        'PUT',
        f'spectrum/{spectrum_id}',
        data={'group_ids': [group_id2]},
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    status, data = api(
        'GET',
        f'spectrum/{spectrum_id}',
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['label'] == custom_label
    group_ids = [g['id'] for g in data['data']['groups']]
    assert group_id1 in group_ids  # PUT is only allowed to remove groups
    assert group_id2 in group_ids
    num_groups = len(data['data']['groups'])

    # adding the same group ID doesn't make redundant groups
    status, data = api(
        'PUT',
        f'spectrum/{spectrum_id}',
        data={'group_ids': [group_id1]},
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    status, data = api(
        'GET',
        f'spectrum/{spectrum_id}',
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    assert num_groups == len(data['data']['groups'])

def test_spectrum_filtering_obj_groups(
    super_admin_user,
    super_admin_token,
    public_source,
    public_source_two_groups,
    lris,
):
    # make groups that must be unique to this test
    status, data = api(
        "POST",
        "groups",
        data={
            "name": str(uuid.uuid4()),
            "group_admins": [super_admin_user.id],
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data["status"] == "success"
    group_id1 = data["data"]["id"]

    status, data = api(
        "POST",
        "groups",
        data={
            "name": str(uuid.uuid4()),
            "group_admins": [super_admin_user.id],
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data["status"] == "success"
    group_id2 = data["data"]["id"]

    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': public_source.id,
            'observed_at': '2020-01-10T00:00:00',
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.3, 232.1, 235.3],
            'group_ids': [group_id1],
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    spectrum_id1 = data['data']['id']

    assert status == 200
    assert data['status'] == 'success'

    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': public_source_two_groups.id,
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [434.7, 432.1, 435.3],
            'group_ids': [group_id1, group_id2],
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    spectrum_id2 = data['data']['id']

    # filter on groups:
    status, data = api(
        'GET',
        'spectra',
        params={'groupIDs': group_id1},  # should get both spectra
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 2
    assert data['data'][0]['id'] == spectrum_id1
    assert data['data'][1]['id'] == spectrum_id2
    assert data['data'][0]['fluxes'][0] == 234.3
    assert data['data'][1]['fluxes'][0] == 434.7

    status, data = api(
        'GET',
        'spectra',
        params={'groupIDs': f'{group_id1}, {group_id2}'},  # should get both spectra
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 2
    assert data['data'][0]['id'] == spectrum_id1
    assert data['data'][1]['id'] == spectrum_id2

    status, data = api(
        'GET',
        'spectra',
        params={'groupIDs': group_id2},  # should get only second spectrum
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id2

    # test objID
    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id1,
            'objID': public_source.id,  # should get only first spectrum
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id1

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id1,
            # partial match to second spectrum
            'objID': public_source_two_groups.id[5:15],
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id2

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id1,
            'objID': 'ZTF2022abcdef',  # should not match anything
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 0

def test_spectrum_filtering_time_ranges(
    super_admin_user,
    super_admin_token,
    public_source,
    lris,
):
    # make a group that is unique to this test
    status, data = api(
        "POST",
        "groups",
        data={
            "name": str(uuid.uuid4()),
            "group_admins": [super_admin_user.id],
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data["status"] == "success"
    group_id = data["data"]["id"]

    # post two spectra at different times
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': public_source.id,
            'observed_at': '2020-01-10T00:00:00',
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.3, 232.1, 235.3],
            'group_ids': [group_id],
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    spectrum_id1 = data['data']['id']

    time_after_posting_first_spec = str(datetime.datetime.utcnow())

    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': public_source.id,
            'observed_at': time_after_posting_first_spec,
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [434.7, 432.1, 435.3],
            'group_ids': [group_id],
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    spectrum_id2 = data['data']['id']

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'observedBefore': '2021-01-10T00:00:00',  # one year after 1st spectrum
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id1
    assert data['data'][0]['fluxes'][0] == 234.3
    assert data['data'][0]['obj_id'] == public_source.id

    # test open ended range that includes second spectrum
    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'observedAfter': time_after_posting_first_spec,
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id2
    assert data['data'][0]['fluxes'][0] == 434.7
    assert data['data'][0]['obj_id'] == public_source.id

    # test open ended range that includes both spectra
    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'observedAfter': '2020-01-01T00:00:00',
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 2

    # test various date formats
    # can't parse this: 'T00:00:00+00:00'
    dates = ['', 'T00:00:00+00:00', 'T00:00:00Z']
    for d in dates:
        status, data = api(
            'GET',
            'spectra',
            params={
                'groupIDs': group_id,
                'observedAfter': f'2020-01-15{d}',  # should get only second spectrum
            },
            token=super_admin_token,
        )

        assert status == 200
        assert data['status'] == 'success'
        assert len(data['data']) == 1
        assert data['data'][0]['id'] == spectrum_id2

def test_spectrum_filtering_id_lists(
    super_admin_user,
    super_admin_token,
    comment_token,
    public_source,
    lris,
    sedm,
    public_source_followup_request,
    public_source_group2_followup_request,
    public_assignment,
):
    # make a group that is unique to this test
    status, data = api(
        "POST",
        "groups",
        data={
            "name": str(uuid.uuid4()),
            "group_admins": [super_admin_user.id],
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data["status"] == "success"
    group_id = data["data"]["id"]

    # post two spectra with very different properties
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': public_source.id,
            'observed_at': '2020-01-10T00:00:00',
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.3, 232.1, 235.3],
            'group_ids': [group_id],
            'followup_request_id': public_source_followup_request.id,
            'assignment_id': public_assignment.id,
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    spectrum_id1 = data['data']['id']

    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': public_source.id,
            'observed_at': str(datetime.datetime.utcnow()),
            'instrument_id': sedm.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [434.7, 432.1, 435.3],
            'group_ids': [group_id],
            'followup_request_id': public_source_group2_followup_request.id,
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    spectrum_id2 = data['data']['id']

    # test instrument IDs
    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'instrumentIDs': lris.id,  # should get only first spectrum
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id1

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'instrumentIDs': sedm.id,  # should get only second spectrum
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id2

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'instrumentIDs': f'{lris.id}, {sedm.id}',  # should get both
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 2

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'instrumentIDs': f'{lris.id}, {sedm.id}, {lris.id * sedm.id}',  # should fail
        },
        token=super_admin_token,
    )

    assert status == 400
    assert data['status'] == 'error'
    assert 'Not all Instrument IDs' in data['message']

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'instrumentIDs': 'free text',  # should fail
        },
        token=super_admin_token,
    )

    assert status == 400
    assert data['status'] == 'error'
    assert 'Could not parse all elements to integers' in data['message']

    # test followup request IDs
    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'followupRequestIDs': public_source_followup_request.id,  # should get only first spectrum
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id1

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'followupRequestIDs': public_source_group2_followup_request.id,  # should only get second spectrum
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id2

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            # should get error
            'followupRequestIDs': public_source_group2_followup_request.id * 10,
        },
        token=super_admin_token,
    )

    assert status == 400
    assert data['status'] == 'error'
    assert 'Not all FollowupRequest IDs' in data['message']

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'followupRequestIDs': public_source_group2_followup_request.id,
        },
        token=comment_token,  # should fail due to permission to see followup request
    )

    assert status == 400
    assert data['status'] == 'error'
    assert 'Not all FollowupRequest IDs' in data['message']

    # test assignments
    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'assignmentIDs': public_assignment.id,  # should get only first spectrum
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id1

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'assignmentIDs': public_assignment.id * 10,  # should fail
        },
        token=super_admin_token,
    )

    assert status == 400
    assert data['status'] == 'error'
    assert 'Not all ClassicalAssignment IDs' in data['message']

def test_spectrum_filtering_origin_label_type(
    super_admin_user,
    super_admin_token,
    public_source,
    lris,
):
    # make a group that is unique to this test
    status, data = api(
        "POST",
        "groups",
        data={
            "name": str(uuid.uuid4()),
            "group_admins": [super_admin_user.id],
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data["status"] == "success"
    group_id = data["data"]["id"]

    # post two spectra with very different properties
    custom_label = str(uuid.uuid4())
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': public_source.id,
            'observed_at': '2020-01-10T00:00:00',
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.3, 232.1, 235.3],
            'group_ids': [group_id],
            'label': custom_label,
            'origin': 'Keck telescope',
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    spectrum_id1 = data['data']['id']

    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': public_source.id,
            'observed_at': str(datetime.datetime.utcnow()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [434.7, 432.1, 435.3],
            'group_ids': [group_id],
            'type': 'host',
            'origin': 'Palomar 60 inch',
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    spectrum_id2 = data['data']['id']

    # test origin
    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'origin': 'Keck',  # should get only first spectrum
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id1

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'origin': ['Gemini', 'VLT'],  # should get nothing
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 0

    # test label
    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'label': custom_label,  # should get only first spectrum
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id1

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'label': ['one', 'two', 'three'],  # should get nothing
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 0

    # test type
    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'type': 'source',  # should get only first spectrum (default type)
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id1

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'type': 'host',  # should get only second spectrum
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id2

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'type': 'host_center',  # should get nothing
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 0

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'type': 'rainbow',  # should get error, (allowed enum)
        },
        token=super_admin_token,
    )

    assert status == 400
    assert data['status'] == 'error'
    assert 'not in list of allowed spectrum types' in data['message']

def test_spectrum_filtering_comments(
    super_admin_user,
    super_admin_token,
    upload_data_token,
    comment_token,
    public_source,
    lris,
):
    # make a group that is unique to this test
    status, data = api(
        "POST",
        "groups",
        data={
            "name": str(uuid.uuid4()),
            # "group_admins": [super_admin_user.id],
        },
        token=upload_data_token,
    )

    assert status == 200
    assert data["status"] == "success"
    group_id = data["data"]["id"]

    # post two spectra with very different properties
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': public_source.id,
            'observed_at': '2020-01-10T00:00:00',
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.3, 232.1, 235.3],
            'group_ids': [group_id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    spectrum_id1 = data['data']['id']

    comment_text = str(uuid.uuid4())
    status, data = api(
        'POST',
        f'spectra/{spectrum_id1}/comments',
        data={'text': comment_text, 'group_ids': [group_id]},
        token=comment_token,
    )

    assert status == 200
    assert data['status'] == 'success'

    time.sleep(2)
    time_after_posting_first_spec = str(datetime.datetime.utcnow())

    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': public_source.id,
            'observed_at': time_after_posting_first_spec,
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [434.7, 432.1, 435.3],
            'group_ids': [group_id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    spectrum_id2 = data['data']['id']

    status, data = api(
        'POST',
        f'spectra/{spectrum_id2}/comments',
        data={
            'text': "looks like Ia.",
            'group_ids': [group_id],
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'commentsFilter': comment_text[10:20],  # should get first spectrum
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id1

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            # should get second spectrum
            'commentsFilterAuthor': super_admin_user.username[8:16],
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id2

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'commentsFilterAuthor': str(uuid.uuid4()),  # should get nothing
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 0
    time_offset = (
        datetime.datetime.utcnow() - datetime.datetime.now()
    ) / datetime.timedelta(hours=1)

    comment_created_time = str(
        arrow.get(time_after_posting_first_spec)
        .shift(seconds=-1)
        .shift(hours=time_offset)
    )
    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'commentsFilterBefore': comment_created_time,  # should get first spectrum
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id1

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'commentsFilterAfter': comment_created_time,  # should get second spectrum
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']) == 1
    assert data['data'][0]['id'] == spectrum_id2

def test_minimal_spectrum(
    super_admin_token,
    public_source,
    lris,
    public_assignment,
    public_source_followup_request,
):
    # make a group that is unique to this test
    status, data = api(
        "POST",
        "groups",
        data={
            "name": str(uuid.uuid4()),
            # "group_admins": [super_admin_user.id],
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data["status"] == "success"
    group_id = data["data"]["id"]

    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': public_source.id,
            'observed_at': '2020-01-10T00:00:00',
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.3, 232.1, 235.3],
            'group_ids': [group_id],
            'followup_request_id': public_source_followup_request.id,
            'assignment_id': public_assignment.id,
            'origin': str(uuid.uuid4()),
            'type': 'host',
            'label': str(uuid.uuid4()),
            'altdata': {'one': 1, 'two': 2},
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    spectrum_id = data['data']['id']

    # post a comment and an annotation as well
    status, data = api(
        'POST',
        f'spectra/{spectrum_id}/comments',
        data={'text': str(uuid.uuid4()), 'group_ids': [group_id]},
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'

    status, data = api(
        'POST',
        f'spectra/{spectrum_id}/annotations',
        data={
            'data': {
                'Gaia_Rp': 14.7,
                'Gaia_Bp': 15.2,
                'Gaia_G': 14.9,
                'period': 13.4,
            },
            'origin': 'Kowalski',
            'group_ids': [group_id],
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'

    status, data = api(
        'GET',
        f'spectra/{spectrum_id}',
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert type(data['data']) == dict
    assert data['data']['id'] == spectrum_id
    single_spec = data['data']

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'minimalPayload': False,
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    assert type(data['data']) == list
    assert len(data['data']) == 1
    full_spec = data['data'][0]

    status, data = api(
        'GET',
        'spectra',
        params={
            'groupIDs': group_id,
            'minimalPayload': True,
        },
        token=super_admin_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    assert type(data['data']) == list
    assert len(data['data']) == 1
    minimal_spec = data['data'][0]

    list_of_keys = [
        'id',
        'altdata',
        'assignment_id',
        'followup_request_id',
        'instrument_id',
        'label',
        'obj_id',
        'observed_at',
        'origin',
        'owner_id',
        'type',
        'original_file_filename',
    ]

    # make sure the minimal list of keys exists in each output
    for k in list_of_keys:
        assert k in minimal_spec  # using multiple spectra, minimal output
        assert k in full_spec  # using multiple spectra, full output
        assert k in single_spec  # using single spectra (should be full always)
        assert minimal_spec[k] == full_spec[k]

    # check that keys of full spec, outside the minimal list, are not included in minimal
    for k in full_spec.keys():
        assert k in list_of_keys or k not in minimal_spec

    # make sure full and single are the same
    for k in single_spec.keys():
        assert k in full_spec
        assert single_spec[k] == full_spec[k]

def test_token_user_get_range_spectrum(
    upload_data_token, public_source, public_group, lris
):
    # post two spectra at two different dates
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': '2020-01-10T00:00:00',
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [434.2, 432.1, 435.3],
            'group_ids': [public_group.id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    # test range that includes first spectrum
    status, data = api(
        'GET',
        f'spectrum/range?instrument_ids={lris.id}'
        f'&min_date=2020-01-01T00:00:00&max_date=2020-01-15T00:00:00',
        token=upload_data_token,
    )
    assert status == 200
    assert len(data['data']) == 1
    assert data['status'] == 'success'
    assert data['data'][0]['fluxes'][0] == 234.2
    assert data['data'][0]['obj_id'] == public_source.id

    # test open ended range that includes second spectrum
    status, data = api(
        'GET',
        f'spectrum/range?instrument_ids={lris.id}&min_date=2020-01-15T00:00:00',
        token=upload_data_token,
    )
    assert status == 200
    assert len(data['data']) == 1
    assert data['status'] == 'success'
    assert data['data'][0]['fluxes'][0] == 434.2
    assert data['data'][0]['obj_id'] == public_source.id

    # test open ended range that includes both spectra
    status, data = api(
        'GET',
        f'spectrum/range?instrument_ids={lris.id}&min_date=2020-01-01T00:00:00',
        token=upload_data_token,
    )
    assert status == 200
    assert len(data['data']) == 2
    assert data['status'] == 'success'

    # test legal variations on input isot format
    # 2020-01-15
    status, data = api(
        'GET',
        f'spectrum/range?instrument_ids={lris.id}&min_date=2020-01-15',
        token=upload_data_token,
    )
    assert status == 200
    assert len(data['data']) == 1
    assert data['status'] == 'success'
    assert data['data'][0]['fluxes'][0] == 434.2
    assert data['data'][0]['obj_id'] == public_source.id

    # 2020-01-15T00:00:00+00:00
    status, data = api(
        'GET',
        f'spectrum/range?instrument_ids={lris.id}&min_date=2020-01-15T00:00:00+00:00',
        token=upload_data_token,
    )
    assert status == 200
    assert len(data['data']) == 1
    assert data['status'] == 'success'
    assert data['data'][0]['fluxes'][0] == 434.2
    assert data['data'][0]['obj_id'] == public_source.id

    # 2020-01-15T00:00:00Z
    status, data = api(
        'GET',
        f'spectrum/range?instrument_ids={lris.id}&min_date=2020-01-15T00:00:00Z',
        token=upload_data_token,
    )
    assert status == 200
    assert len(data['data']) == 1
    assert data['status'] == 'success'
    assert data['data'][0]['fluxes'][0] == 434.2
    assert data['data'][0]['obj_id'] == public_source.id

    # test with no instrument ids
    status, data = api(
        'GET',
        'spectrum/range?min_date=2020-01-01T00:00:00&max_date=2020-02-01',
        token=upload_data_token,
    )
    assert status == 200
    assert len(data['data']) == 1
    assert data['status'] == 'success'
    assert data['data'][0]['fluxes'][0] == 234.2
    assert data['data'][0]['obj_id'] == public_source.id

def test_token_user_post_get_spectrum_data(
    upload_data_token, public_source, public_group, lris
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    spectrum_id = data['data']['id']
    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['fluxes'][0] == 234.2
    assert data['data']['obj_id'] == public_source.id

def test_token_user_post_spectrum_no_instrument_id(
    upload_data_token, public_source, public_group
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
        },
        token=upload_data_token,
    )
    assert status == 400
    assert data['status'] == 'error'

    # should be a marshamallow error, not a psycopg2 error
    # (see https://github.com/skyportal/skyportal/issues/1047)
    assert 'psycopg2' not in data['message']

def test_token_user_post_spectrum_all_groups(
    upload_data_token, public_source_two_groups, lris
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source_two_groups.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': "all",
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    spectrum_id = data['data']['id']
    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['fluxes'][0] == 234.2
    assert data['data']['obj_id'] == public_source_two_groups.id

def test_token_user_post_spectrum_no_access(
    view_only_token, public_source, public_group, lris
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
        },
        token=view_only_token,
    )
    assert status == 401
    assert data['status'] == 'error'

def test_token_user_update_spectrum(
    upload_data_token, public_source, public_group, lris
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    spectrum_id = data['data']['id']
    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['fluxes'][0] == 234.2

    status, data = api(
        'PUT',
        f'spectrum/{spectrum_id}',
        data={
            'fluxes': [222.2, 232.1, 235.3],
            'observed_at': str(datetime.datetime.now()),
            'wavelengths': [664, 665, 666],
            'group_ids': 'all',
        },
        token=upload_data_token,
    )

    assert status == 200
    assert data['status'] == 'success'

    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['fluxes'][0] == 222.2
    # test that length of groups is greater than 1 after adding all groups to the spectrum
    assert len(data['data']['groups']) > 1

def test_token_user_cannot_update_unowned_spectrum(
    upload_data_token, manage_sources_token, public_source, public_group, lris
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    spectrum_id = data['data']['id']
    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['fluxes'][0] == 234.2

    status, data = api(
        'PUT',
        f'spectrum/{spectrum_id}',
        data={
            'fluxes': [222.2, 232.1, 235.3],
            'observed_at': str(datetime.datetime.now()),
            'wavelengths': [664, 665, 666],
        },
        token=manage_sources_token,
    )

    assert status == 401
    assert data['status'] == 'error'

def test_admin_can_update_unowned_spectrum_data(
    upload_data_token, super_admin_token, public_source, public_group, lris
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    spectrum_id = data['data']['id']
    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['fluxes'][0] == 234.2

    status, data = api(
        'PUT',
        f'spectrum/{spectrum_id}',
        data={
            'fluxes': [222.2, 232.1, 235.3],
            'observed_at': str(datetime.datetime.now()),
            'wavelengths': [664, 665, 666],
            'group_ids': [2, 3],
        },
        token=super_admin_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['fluxes'][0] == 222.2
    # check if length of groups is 4 after adding permission to two groups (groups with id 2 and 3) because two groups already have permission to this spectrum (groups with id 1405 and 1406)
    assert len(data['data']['groups']) == 4

def test_spectrum_owner_id_is_unmodifiable(
    upload_data_token,
    super_admin_user,
    super_admin_token,
    public_source,
    public_group,
    lris,
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    spectrum_id = data['data']['id']
    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['fluxes'][0] == 234.2

    status, data = api(
        'PUT',
        f'spectrum/{spectrum_id}',
        data={'owner_id': super_admin_user.id},
        token=super_admin_token,
    )

    assert status == 400
    assert data['status'] == 'error'

def test_user_cannot_delete_unowned_spectrum_data(
    upload_data_token, manage_sources_token, public_source, public_group, lris
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    spectrum_id = data['data']['id']
    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['fluxes'][0] == 234.2
    assert data['data']['obj_id'] == public_source.id

    status, data = api('DELETE', f'spectrum/{spectrum_id}', token=manage_sources_token)
    assert status == 401

def test_user_can_delete_owned_spectrum_data(
    upload_data_token, public_source, public_group, lris
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    spectrum_id = data['data']['id']
    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['fluxes'][0] == 234.2
    assert data['data']['obj_id'] == public_source.id

    status, data = api('DELETE', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200

    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 403

def test_admin_can_delete_unowned_spectrum_data(
    upload_data_token, super_admin_token, public_source, public_group, lris
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    spectrum_id = data['data']['id']
    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['fluxes'][0] == 234.2
    assert data['data']['obj_id'] == public_source.id

    status, data = api('DELETE', f'spectrum/{spectrum_id}', token=super_admin_token)
    assert status == 200

    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 403

def test_jsonify_spectrum_header(
    upload_data_token, manage_sources_token, public_source, public_group, lris
):
    for filename in glob(f'{os.path.dirname(__file__)}/../data/ZTF*.ascii.head'):
        with open(filename[:-5]) as f:
            status, data = api(
                'POST',
                'spectrum/parse/ascii',
                data={
                    'fluxerr_column': 3
                    if 'ZTF20abpuxna_20200915_Keck1_v1.ascii' in filename
                    else 2
                    if 'P60' in filename
                    else None,
                    'ascii': f.read(),
                },
                token=upload_data_token,
            )
        assert status == 200
        assert data['status'] == 'success'

        answer = yaml.load(open(filename), Loader=yaml.FullLoader)

        # check the header serialization
        for key in answer:
            # special keys
            if key not in ['COMMENT', 'END', 'HISTORY']:
                if isinstance(data['data']['altdata'][key], dict):
                    value = data['data']['altdata'][key]['value']
                else:
                    value = data['data']['altdata'][key]
                if isinstance(answer[key], (str, int)):
                    assert str(value) == str(answer[key])
                elif isinstance(answer[key], datetime.datetime):
                    assert datetime.datetime.fromisoformat(value) == answer[key]
                elif isinstance(answer[key], datetime.date):
                    assert datetime.datetime.fromisoformat(value).date() == answer[key]
                elif answer[key] is None:
                    assert value is None
                else:
                    np.testing.assert_allclose(value, answer[key])

def test_can_post_spectrum_no_groups(
    upload_data_token, public_source, public_group, lris
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    spectrum_id = data['data']['id']
    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']['groups']) == 1

def test_can_post_spectrum_empty_groups_list(
    upload_data_token, public_source, public_group, lris
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    spectrum_id = data['data']['id']
    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert len(data['data']['groups']) == 1

def test_jsonify_spectrum_data(
    upload_data_token, manage_sources_token, public_source, public_group, lris
):
    for filename in glob(f'{os.path.dirname(__file__)}/../data/ZTF*.ascii'):
        with open(filename) as f:
            status, data = api(
                'POST',
                'spectrum/parse/ascii',
                data={
                    'fluxerr_column': 3
                    if 'ZTF20abpuxna_20200915_Keck1_v1.ascii' in filename
                    else 2
                    if 'P60' in filename
                    else None,
                    'ascii': f.read(),
                },
                token=upload_data_token,
            )
        assert status == 200
        assert data['status'] == 'success'

        answer = np.genfromtxt(filename, dtype=float, encoding='ascii')

        if answer.shape[-1] == 2:
            np.testing.assert_allclose(
                np.asarray(data['data']['wavelengths'], dtype=float), answer[:, 0]
            )
            np.testing.assert_allclose(
                np.asarray(data['data']['fluxes'], dtype=float), answer[:, 1]
            )

        elif answer.shape[-1] == 3:
            np.testing.assert_allclose(
                np.asarray(data['data']['wavelengths'], dtype=float), answer[:, 0]
            )
            np.testing.assert_allclose(
                np.asarray(data['data']['fluxes'], dtype=float), answer[:, 1]
            )
            np.testing.assert_allclose(
                np.asarray(data['data']['errors'], dtype=float), answer[:, 2]
            )

        else:
            # this is the long one from Keck
            np.testing.assert_allclose(
                np.asarray(data['data']['wavelengths'], dtype=float), answer[:, 0]
            )
            np.testing.assert_allclose(
                np.asarray(data['data']['fluxes'], dtype=float), answer[:, 1]
            )
            np.testing.assert_allclose(
                np.asarray(data['data']['errors'], dtype=float), answer[:, 3]
            )

def test_upload_bad_spectrum_from_ascii_file(
    upload_data_token, manage_sources_token, public_source, public_group, lris
):
    for filename in glob(f'{os.path.dirname(__file__)}/../data/ZTF*.ascii.bad'):
        with open(filename) as f:
            content = f.read()
            observed_at = str(datetime.datetime.now())

            status, data = api(
                'POST',
                'spectrum/ascii',
                data={
                    'obj_id': str(public_source.id),
                    'observed_at': observed_at,
                    'instrument_id': lris.id,
                    'group_ids': [public_group.id],
                    'fluxerr_column': 3
                    if 'ZTF20abpuxna_20200915_Keck1_v1.ascii' in filename
                    else 2
                    if 'P60' in filename
                    else None,
                    'ascii': content,
                    'filename': filename,
                },
                token=upload_data_token,
            )

            assert status == 400
            assert data['status'] == 'error'

def test_token_user_post_to_foreign_group_and_retrieve(
    upload_data_token, public_source_two_groups, public_group2, lris
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source_two_groups.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group2.id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    spectrum_id = data['data']['id']
    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200

def test_parse_integer_spectrum_ascii(upload_data_token):

    status, data = api(
        'POST',
        'spectrum/parse/ascii',
        data={'ascii': '4000 0.01\n4500 0.02\n5000 0.005\n5500 0.006\n6000 0.01\n'},
        token=upload_data_token,
    )

    assert status == 200
    assert data['status'] == 'success'

    for wave in data['data']['wavelengths']:
        assert isinstance(wave, float)

def test_spectrum_external_reducer_and_observer(
    upload_data_token, public_source, public_group, lris, user
):
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
            'reduced_by': [user.id],
            'external_reducer': "Test external reducer",
            'observed_by': [user.id],
            'external_observer': "Test external observer",
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'

    spectrum_id = data['data']['id']
    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['reducers'][0]['id'] == user.id
    assert data['data']['observers'][0]['id'] == user.id
    assert data['data']['external_reducer'] == "Test external reducer"
    assert data['data']['external_observer'] == "Test external observer"

def test_spectrum_plot_only_some_spectra(
    upload_data_token, public_source, public_group, lris, sedm
):
    # upload spectrum 1
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    spec_id1 = data['data']['id']

    # upload spectrum 2
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': sedm.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [134.2, 132.1, 135.3],
            'group_ids': [public_group.id],
        },
        token=upload_data_token,
    )

    assert status == 200
    assert data['status'] == 'success'
    spec_id2 = data['data']['id']

    status, data = api(
        'GET', f'internal/plot/spectroscopy/{public_source.id}', token=upload_data_token
    )

    assert status == 200

    # check how many different instruments are plotted
    instruments = set()
    refs = data['data']['bokehJSON']['doc']['roots']['references']
    for r in refs:
        if 'data' in r['attributes'] and 'instrument' in r['attributes']['data']:
            instruments.add(r['attributes']['data']['instrument'][0])

    assert len(instruments) > 2

    # now ask for the same plot but with only one spectrum
    status, data = api(
        'GET',
        f'internal/plot/spectroscopy/{public_source.id}',
        params={'spectrumID': f'{spec_id1},{spec_id2}'},
        token=upload_data_token,
    )

    assert status == 200

    # check how many different instruments are plotted
    instruments = set()
    refs = data['data']['bokehJSON']['doc']['roots']['references']
    for r in refs:
        if 'data' in r['attributes'] and 'instrument' in r['attributes']['data']:
            instruments.add(r['attributes']['data']['instrument'][0])

    assert len(instruments) == 2
    assert 'LRIS' in str(instruments)
    assert 'SEDM' in str(instruments)

def test_post_get_spectrum_type(upload_data_token, public_source, public_group, lris):

    # post this spectrum without a type (should default to "source")
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
        },
        token=upload_data_token,
    )
    assert status == 200
    assert data['status'] == 'success'
    spectrum_id = data['data']['id']

    status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
    assert status == 200
    assert data['status'] == 'success'
    assert data['data']['type'] == default_spectrum_type

    assert default_spectrum_type in ALLOWED_SPECTRUM_TYPES

    if len(ALLOWED_SPECTRUM_TYPES) > 1:

        new_allowed_types = list(ALLOWED_SPECTRUM_TYPES)
        new_allowed_types.remove(default_spectrum_type)

        status, data = api(
            'POST',
            'spectrum',
            data={
                'obj_id': str(public_source.id),
                'observed_at': str(datetime.datetime.now()),
                'instrument_id': lris.id,
                'wavelengths': [664, 665, 666],
                'fluxes': [234.2, 232.1, 235.3],
                'group_ids': [public_group.id],
                'type': new_allowed_types[0],
            },
            token=upload_data_token,
        )
        assert status == 200
        assert data['status'] == 'success'
        spectrum_id = data['data']['id']

        status, data = api('GET', f'spectrum/{spectrum_id}', token=upload_data_token)
        assert status == 200
        assert data['status'] == 'success'
        assert data['data']['type'] == new_allowed_types[0]

def test_post_wrong_spectrum_type(upload_data_token, public_source, public_group, lris):

    # post this spectrum with the wrong type
    status, data = api(
        'POST',
        'spectrum',
        data={
            'obj_id': str(public_source.id),
            'observed_at': str(datetime.datetime.now()),
            'instrument_id': lris.id,
            'wavelengths': [664, 665, 666],
            'fluxes': [234.2, 232.1, 235.3],
            'group_ids': [public_group.id],
            'type': str(uuid.uuid4()),
        },
        token=upload_data_token,
    )
    assert status == 400
    assert 'Must be one of: ' in data['message']

max_stars_repo_name: jadalilleboe/skyportal

max_issues_repo_name: jadalilleboe/skyportal

max_forks_repo_name: jadalilleboe/skyportal


text:
import numpy as np
import matplotlib.pyplot as plt

# constants
hh, cc, kk = 6.6269655e-27, 2.99792458e10, 1.380658e-16

# rough frequencies
nu0 = cc / 0.13
nu1 = cc / 0.9

# temperatures
T = np.linspace(15., 60., 100)

# alpha_Pl
alp_pl = np.zeros_like(T)
for i in range(len(T)):
    foo = np.log((np.exp(hh*nu1/(kk*T[i]))-1.)/(np.exp(hh*nu0/(kk*T[i]))-1.))
    alp_pl[i] = 3. + foo / np.log(nu0/nu1)

plt.plot(T, alp_pl)
plt.show()

ext: py

max_stars_repo_name: seanandrews/ARAA

max_issues_repo_name: seanandrews/ARAA

max_forks_repo_name: seanandrews/ARAA

reason: import numpy


text:
import json
import os
from konlpy.tag import Okt
import nltk
import numpy as np

from tensorflow.keras import models
from tensorflow.keras import layers
from tensorflow.keras import optimizers
from tensorflow.keras import losses
from tensorflow.keras import metrics

okt = Okt()

def tokenize(doc):
    return ['/'.join(t) for t in okt.pos(doc, norm=True, stem=True)]

def data_load(filename):
    with open(filename, 'r', encoding='UTF8') as f:
        data = [line.split('\t') for line in f.read().splitlines()]
        data = data[1:]
    return data

def term_frequency(doc):
    return [doc.count(word) for word in selected_words]

train_data = data_load('ratings_train_min.txt')
test_data = data_load('ratings_test_min.txt')

print("전체 데이터 갯수 : {}".format(len(train_data)))
print("테스트용 데이터 갯수 : {}".format(len(test_data)))

if os.path.isfile('train_docs.json'):
    with open('train_docs.json', encoding='UTF8') as f:
        train_docs = json.load(f)
    with open('test_docs.json', encoding='UTF8') as f:
        test_docs = json.load(f)
else:
    train_docs = [(tokenize(row[1]), row[2]) for row in train_data]
    test_docs = [(tokenize(row[1]), row[2]) for row in test_data]
    with open('train_docs.json', 'w', encoding='UTF8') as make_file:
        json.dump(train_docs, make_file, ensure_ascii=False, indent="\t")
    with open('test_docs.json', 'w', encoding='UTF8') as make_file:
        json.dump(test_docs, make_file, ensure_ascii=False, indent="\t")

print("문장 태깅 결과 (5개만 보기) : ")
for index in range(5):
    print(train_docs[index])

tokens = [t for d in train_docs for t in d[0]]
print("토큰 갯수 : ")
print(len(tokens))

text = nltk.Text(tokens, name='NMSC')

print("전체 토큰 : {}".format(len(text.tokens)))
print("중복 제외 토큰 : {}".format(len(set(text.tokens))))
print("출현빈도 상위 TOP 10 토큰 : {}".format(text.vocab().most_common(10)))

selected_words = [f[0] for f in text.vocab().most_common(1000)]

train_x = [term_frequency(d) for d, _ in train_docs]
test_x = [term_frequency(d) for d, _ in test_docs]
train_y = [c for _, c in train_docs]
test_y = [c for _, c in test_docs]

x_train = np.asarray(train_x).astype('float32')
x_test = np.asarray(test_x).astype('float32')

y_train = np.asarray(train_y).astype('float32')
y_test = np.asarray(test_y).astype('float32')

model = models.Sequential()
model.add(layers.Dense(64, activation='relu', input_shape=(1000,)))
model.add(layers.Dense(64, activation='relu'))
model.add(layers.Dense(1, activation='sigmoid'))

model.compile(optimizer=optimizers.RMSprop(lr=0.001), loss=losses.binary_crossentropy, metrics=[metrics.binary_accuracy])

model.fit(x_train, y_train, epochs=10, batch_size=512)
results = model.evaluate(x_test, y_test)

print("정확도 : {}".format(results))

## 학습 모델 결과 저장
model.save('./result_model.mod')

## selected_words 리스트 저장
with open('selected_words.list', 'w') as file:
    for index in range(0, len(selected_words)):
        file.writelines(selected_words[index])
        file.writelines('\n')


text:
import argparse
import os
import sys

import numpy as np

# to make run from console for module import
sys.path.append(os.path.abspath('..'))

from main.config import Config
from main.model import Model
from visualise.trimesh_renderer import TrimeshRenderer
from visualise.vis_util import preprocess_image, visualize

def str2bool(v):
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Demo HMR2.0')

    parser.add_argument('--image', required=False, default='coco1.png')
    parser.add_argument('--model', required=False, default='base_model', help="model from logs folder")
    parser.add_argument('--setting', required=False, default='paired(joints)', help="setting of the model")
    parser.add_argument('--joint_type', required=False, default='cocoplus', help="<cocoplus|custom>")
    parser.add_argument('--init_toes', required=False, default=False, type=str2bool,
                        help="only set to True when joint_type=cocoplus")

    args = parser.parse_args()
    if args.init_toes:
        assert args.joint_type, "Only init toes when joint type is cocoplus!"

    class DemoConfig(Config):
        BATCH_SIZE = 1
        ENCODER_ONLY = True
        LOG_DIR = os.path.abspath('../../logs/{}/{}'.format(args.setting, args.model))
        INITIALIZE_CUSTOM_REGRESSOR = args.init_toes
        JOINT_TYPE = args.joint_type

    config = DemoConfig()

    # initialize model
    model = Model()
    original_img, input_img, params = preprocess_image('images/{}'.format(args.image), config.ENCODER_INPUT_SHAPE[0])

    result = model.detect(input_img)

    cam = np.squeeze(result['cam'].numpy())[:3]
    vertices = np.squeeze(result['vertices'].numpy())
    joints = np.squeeze(result['kp2d'].numpy())
    joints = ((joints + 1) * 0.5) * params['img_size']

    renderer = TrimeshRenderer()
    visualize(renderer, original_img, params, vertices, cam, joints)


text:
# Copyright 2022 Google LLC.
#
# 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.

"""Class for relative position biases.

T5 uses a form of relative attention which biases the attention matrix, so each
head effectively attends to things at different scales, irrespective of the
contents of keys and queries.

In the future, this class may be unified with classes which take into account
key and query contents, like the original relative position embeddings of Shaw
et al. and new proposals. However, this will rely on XLA to recover efficiency
for this class (especially when, as in the original T5, the same bias matrix
is shared for all layers).
"""
from typing import Any, Callable

from flax import linen as nn
from flax.linen import partitioning
from jax import lax
import jax.numpy as jnp
import numpy as np

from flaxformer.types import Array

class RelativePositionBiases(nn.Module):
  """Adds T5-style relative positional embeddings to the attention logits.

  Attributes:
    num_buckets: Number of buckets to bucket distances between key and query
      positions into.
    max_distance: Maximum distance before everything is lumped into the last
      distance bucket.
    num_heads: Number of heads in the attention layer. Each head will get a
      different relative position weighting.
    dtype: Type of arrays through this module.
    embedding_init: initializer for relative embedding table.
  """
  num_buckets: int
  max_distance: int
  num_heads: int
  dtype: Any
  embedding_init: Callable[..., Array] = nn.linear.default_embed_init

  @staticmethod
  def _relative_position_bucket(relative_position,
                                bidirectional=True,
                                num_buckets=32,
                                max_distance=128):
    """Translate relative position to a bucket number for relative attention.

    The relative position is defined as memory_position - query_position, i.e.
    the distance in tokens from the attending position to the attended-to
    position.  If bidirectional=False, then positive relative positions are
    invalid.
    We use smaller buckets for small absolute relative_position and larger
    buckets for larger absolute relative_positions.  All relative
    positions >=max_distance  map to the same bucket.  All relative
    positions <=-max_distance map to the same bucket.  This should allow for
    more graceful generalization to longer sequences than the model has been
    trained on.

    Args:
      relative_position: an int32 array
      bidirectional: a boolean - whether the attention is bidirectional
      num_buckets: an integer
      max_distance: an integer

    Returns:
      a Tensor with the same shape as relative_position, containing int32
        values in the range [0, num_buckets)
    """
    ret = 0
    n = -relative_position
    if bidirectional:
      num_buckets //= 2
      ret += (n < 0).astype(np.int32) * num_buckets
      n = np.abs(n)
    else:
      n = np.maximum(n, 0)
    # now n is in the range [0, inf)
    max_exact = num_buckets // 2
    is_small = (n < max_exact)
    val_if_large = max_exact + (
        np.log(n.astype(np.float32) / max_exact + np.finfo(np.float32).eps) /
        np.log(max_distance / max_exact) *
        (num_buckets - max_exact)).astype(np.int32)
    val_if_large = np.minimum(val_if_large, num_buckets - 1)
    ret += np.where(is_small, n, val_if_large)
    return ret

  @nn.compact
  def __call__(self, qlen, klen, bidirectional=True, decode=False):
    """Produce relative position embedding attention biases.

    Args:
      qlen: attention query length.
      klen: attention key length.
      bidirectional: whether to allow positive memory-query relative position
        embeddings.
      decode: whether to cache relative position bias during autoregressive
        decoding.

    Returns:
      output: `(1, num_heads, q_len, k_len)` attention bias
    """
    # bidirectional embeddings don't make sense when decoding (and break cache).
    if decode and bidirectional:
      raise ValueError(
          'bidirectional RelativePositionBiases are not supported when decode=True.'
      )

    # We only cache the bias if the model was already initialized, i.e. if this
    # module is called with model.apply and decode = True. We raise an error if
    # called with model.init and decode = True, since this can cache incorrect
    # positional embeddings produced by random parameters.
    is_initialized = self.has_variable('params', 'rel_embedding')
    if decode and not is_initialized:
      raise ValueError(
          'decode-mode cannot be enabled during init. use model.apply to '
          'initialize the decoding cache.')

    # Return pre-computed relative position bias in cache during decode steps.
    if decode and self.has_variable('cache', 'cached_bias'):
      cached_bias = self.get_variable('cache', 'cached_bias')
      expected_bias_shape = (1, self.num_heads, qlen, klen)
      if cached_bias.shape != expected_bias_shape:
        raise ValueError(f'The cached relative position attention bias was '
                         f'expected to have shape {expected_bias_shape} but '
                         f'instead has the shape {cached_bias.shape}.')
      return cached_bias

    # TODO: should we be computing this w. numpy as a program
    # constant?
    context_position = np.arange(qlen, dtype=jnp.int32)[:, None]
    memory_position = np.arange(klen, dtype=jnp.int32)[None, :]
    relative_position = memory_position - context_position  # shape (qlen, klen)
    rp_bucket = self._relative_position_bucket(
        relative_position,
        bidirectional=bidirectional,
        num_buckets=self.num_buckets,
        max_distance=self.max_distance)
    relative_attention_bias = partitioning.param_with_axes(
        'rel_embedding',
        self.embedding_init, (self.num_heads, self.num_buckets),
        jnp.float32,
        axes=('heads', 'relpos_buckets'))
    relative_attention_bias = jnp.asarray(relative_attention_bias, self.dtype)
    # Instead of using a slow gather, we create a leading-dimension one-hot
    # array from rp_bucket and use it to perform the gather-equivalent via a
    # contraction, i.e.:
    # (num_head, num_buckets) x (num_buckets one-hot, qlen, klen).
    # This is equivalent to relative_attention_bias[:, rp_bucket]
    bcast_iota = lax.broadcasted_iota(jnp.int32, (self.num_buckets, 1, 1), 0)
    rp_bucket_one_hot = jnp.array(
        rp_bucket[jnp.newaxis, ...] == bcast_iota, dtype=self.dtype)
    # --> shape (qlen, klen, num_heads)
    values = lax.dot_general(
        relative_attention_bias,
        rp_bucket_one_hot,
        (
            ((1,), (0,)),  # rhs, lhs contracting dims
            ((), ())))  # no batched dims
    # Add a singleton batch dimension.
    # --> shape (1, num_heads, qlen, klen)
    out = values[jnp.newaxis, ...]

    # Store computed relative position bias in cache after first calculation.
    if decode:
      _ = self.variable('cache', 'cached_bias', lambda: out)

    return out

max_stars_repo_name: tomweingarten/flaxformer

max_issues_repo_name: tomweingarten/flaxformer

max_forks_repo_name: tomweingarten/flaxformer

reason: import numpy,import jax,from jax


text:
"""
======================================
Predict age from white matter features
======================================

Predict subject age from white matter features. This example fetches the
Weston-Havens dataset described in Yeatman et al [1]_. This dataset contains
tractometry features from 77 subjects ages 6-50. The plots display the absolute
value of the mean regression coefficients (averaged across cross-validation
splits) for the mean diffusivity (MD) features.

Predictive performance for this example is quite poor. In a research setting,
one might have to ensemble a number of SGL estimators together and conduct a
more thorough search of the hyperparameter space. For more details, please see
[2]_.

.. [1]  Jason D. Yeatman, Brian A. Wandell, & Aviv A. Mezer, "Lifespan
    maturation and degeneration of human brain white matter" Nature
    Communications, vol. 5:1, pp. 4932, 2014 DOI: 10.1038/ncomms5932

.. [2]  Adam Richie-Halford, Jason Yeatman, Noah Simon, and Ariel Rokem
   "Multidimensional analysis and detection of informative features in human
   brain white matter" PLOS Computational Biology, 2021 DOI:
   10.1371/journal.pcbi.1009136

"""
import matplotlib.pyplot as plt
import numpy as np
import os.path as op

from afqinsight.datasets import download_weston_havens, load_afq_data
from afqinsight import make_afq_regressor_pipeline

from sklearn.model_selection import cross_validate

##########################################################################
# Fetch example data
# ------------------
#
# The :func:`download_weston_havens` function download the data used in this
# example and places it in the `~/.cache/afq-insight/weston_havens` directory.
# If the directory does not exist, it is created. The data follows the format
# expected by the :func:`load_afq_data` function: a file called `nodes.csv` that
# contains AFQ tract profiles and a file called `subjects.csv` that contains
# information about the subjects. The two files are linked through the
# `subjectID` column that should exist in both of them. For more information
# about this format, see also the `AFQ-Browser documentation
# <https://yeatmanlab.github.io/AFQ-Browser/dataformat.html>`_ (items 2 and 3).

workdir = download_weston_havens()

##########################################################################
# Read in the data
# ----------------
# Next, we read in the data. The :func:`load_afq_data` function expects a string
# input that points to a directory that holds appropriately-shaped data and
# returns variables that we will use below in our analysis of the data.

afqdata = load_afq_data(
    fn_nodes=op.join(workdir, "nodes.csv"),
    fn_subjects=op.join(workdir, "subjects.csv"),
    dwi_metrics=["md", "fa"],
    target_cols=["Age"],
)

# afqdata is a namedtuple. You can access it's fields using dot notation or by
# unpacking the tuple. To see all of the available fields use `afqdata._fields`
X = afqdata.X
y = afqdata.y
groups = afqdata.groups
feature_names = afqdata.feature_names
group_names = afqdata.group_names
subjects = afqdata.subjects

##########################################################################
# Create an analysis pipeline
# ---------------------------
#

pipe = make_afq_regressor_pipeline(
    imputer_kwargs={"strategy": "median"},  # Use median imputation
    use_cv_estimator=True,  # Automatically determine the best hyperparameters
    scaler="standard",  # Standard scale the features before regression
    groups=groups,
    verbose=0,  # Be quiet!
    pipeline_verbosity=False,  # No really, be quiet!
    tuning_strategy="bayes",  # Use BayesSearchCV to determine the optimal hyperparameters
    n_bayes_iter=10,  # Consider only this many points in hyperparameter space
    cv=3,  # Use three CV splits to evaluate each hyperparameter combination
    l1_ratio=[0.0, 1.0],  # Explore the entire range of ``l1_ratio``
    eps=5e-2,  # This is the ratio of the smallest to largest ``alpha`` value
    tol=1e-2,  # Set a lenient convergence tolerance just for this example
)

# ``pipe`` is a scikit-learn pipeline and can be used in other scikit-learn
# functions. For example, here we are doing 5-fold cross-validation using scikit
# learn's :func:`cross_validate` function.

scores = cross_validate(
    pipe, X, y, cv=5, return_train_score=True, return_estimator=True
)

print(f"Mean train score: {np.mean(scores['train_score']):5.3f}")
print(f"Mean test score:  {np.mean(scores['test_score']):5.3f}")
print(f"Mean fit time:    {np.mean(scores['fit_time']):5.2f}s")
print(f"Mean score time:  {np.mean(scores['score_time']):5.2f}s")

mean_coefs = np.mean(
    np.abs([est.named_steps["estimate"].coef_ for est in scores["estimator"]]), axis=0
)

fig, ax = plt.subplots(1, 1, figsize=(8, 8))
_ = ax.plot(mean_coefs[1800:], color="black", lw=2)
_ = ax.set_xlim(0, 1800)

colors = plt.get_cmap("tab20").colors
for grp, grp_name, color in zip(groups[:18], group_names[18:], colors):
    _ = ax.axvspan(grp.min(), grp.max() + 1, color=color, alpha=0.8, label=grp_name[1])

box = ax.get_position()
ax.set_position([box.x0, box.y0 + box.height * 0.375, box.width, box.height * 0.625])

_ = ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.125), ncol=3)
_ = ax.set_ylabel(r"$\hat{\beta}$", fontsize=16)
_ = ax.set_xlabel("Group principal component", fontsize=16)
_ = ax.set_title("Group Principal Regression Coefficients (MD only)", fontsize=18)

max_stars_repo_name: arokem/AFQ-Insight

max_issues_repo_name: arokem/AFQ-Insight

max_forks_repo_name: arokem/AFQ-Insight


text:
'''
This is a test module to test 911CallsStatisicalAnalysisOverTime module's functionality
'''

import calls_statistical_analysis_over_time as ca
import numpy as np
import pandas as pd

#test data clean process functionality.
def test_dataClean():
    '''exception test to check for nan values.
       The function takes in a dataframe that has some columns with missing values
       or Nan values as an argument and then tries to throw a ValueError. If a
       ValueError is caught, the output is pass
    '''
    data1 = pd.read_csv("C:/Users/Runbang Tang/Downloads/Road_Weather_Information_Stations.csv",nrows=500)

    try:
        ca.dataClean(data1)
    except ValueError:
        pass

#exception test to Check that if after data cleaning, the data are in the correct form.
def test_rightForm():
    '''exception test to Check if data cleaning process works properly. If TyprError
    is caught, the output is pass
    '''
    data2 = pd.read_csv("C:/Users/Runbang Tang/Downloads/Seattle_Real_Time_Fire_911_Calls2.csv",nrows=1000)
    df2 = ca.dataProcess(data2)
    try:
        ca.dataRight(df2)
    except TypeError:
        pass