import numpy as np
import cooler
import click
from tqdm import tqdm
import pandas as pd
import numba
from numba import jit
import sys
from sklearn.neighbors import KDTree
import functools
from importlib_resources import files
from matplotlib import  pylab as plt
from scipy.stats import pearsonr as pcc

looptpl = np.loadtxt(str(files('robustad').joinpath('data/looptpl.txt')))

def filterTAD(data, offset, left, right, minRatio=1.1):
    scores = []
    cscore=0
    idx = []
    for i in range(-5, 6):
        _left = left + i
        _right = right + i
        if _left-offset<0 or _right-offset>=data.shape[0]:
            continue
        score = Delta(data=data, offset=offset, left=_left, right=_right, minRatio=minRatio,mask=None)
        if i==0:
            cscore = score
        scores.append(score)
        idx.append(i)
    return scores, idx,cscore

def popOneL0TAD(tads):
    for tad in tads:
        nestTADs = []
        for tad2 in tads:
            if tad2[0] >= tad[0] and tad2[1] <= tad[1]:
                if tad2[0] == tad[0] and tad2[1] == tad[1]:
                    pass
                else:
                    nestTADs.append(tad2)

        if len(nestTADs) == 0:
            tads.remove(tad)
            return tad, tads

def cleanTAD(tads=None,data=None):
    goodtads = []
    badtads = []
    diagExp = []
    for i in range(data.shape[0]):
        diagExp.append(np.nanmean(np.diag(data, i)))
    diagExp = np.asarray(diagExp)+eps

    data[np.diag_indices(data.shape[0], 2)] = np.nan

    while len(tads) > 0:
        tad, tads = popOneL0TAD(tads)
        l, r = tad
        start = max([0, 2 * l - r - 1 - 30])
        end = 2 * r - l + 1 + 30
        m =data[start:end, start:end].copy()
        s, idx,cscore = filterTAD(m, start, l, r)
        peak = np.abs(idx[np.argmax(s)])
        rank = np.argwhere(np.argsort(-np.asarray(s)) == 5).flatten()[0]

        corner = o2oe(data[l - 8:l + 9, r - 8:r + 9], l - 8, r - 8, diagExp)
        if np.sum(~np.isnan(corner)) > 100 and corner.shape == (17,17):
            cornerpcc = pcc(corner[~np.isnan(corner)], looptpl[~np.isnan(corner)])[0]
        else:
            cornerpcc = -1
        offset = 0
        if (peak <= 4 and rank <= 5) or r - l < 20 or cornerpcc > 0.3:
            data[np.max(np.asarray([0, l - offset])):r - offset + 1,
            np.max(np.asarray([0, l - offset])):r - offset + 1] = np.nan
            data[np.max(np.asarray([0, l - offset - 4])):l - offset + 5,
            np.max(np.asarray([0, r - offset - 4])):r - offset + 5] = np.nan  # mask dot corner
            goodtads.append((l, r))
        else:
            badtads.append((l, r))

    return goodtads,badtads

def o2oe(o,l,r,diagExp):
    d=o*0
    for i in range(o.shape[0]):
        for j in range(o.shape[1]):
            d[i,j] = diagExp[abs(r-l + j-i)]
    return o/d

def addback(left,right,data,minTAD,maxTAD):
    minTAD=int(minTAD)
    maxTAD =int(maxTAD)
    result =[]
    diagExp = []
    badnum=0
    left = list(left)
    right = list(right)

    for i in range(data.shape[0]):
        diagExp.append(np.nanmean(np.diag(data, i)))
    diagExp = np.asarray(diagExp)+eps

    maxij=max([max(left),max(right)])
    label=np.zeros(maxij+1)
    for i in range(len(left)):
        label[left[i]] = 1
    for i in range(len(right)):
        if label[right[i]] == 1:
            label[right[i]] = 3
        else:
            label[right[i]]=2
    leftidx = np.argwhere((label==1) | (label==3)).flatten()
    for l in leftidx:
        for r in range(l+minTAD,min([maxij,l+maxTAD+1])):
            if label[r] >= 2:
                corner = o2oe(data[l - 8:l + 9, r - 8:r + 9], l - 8, r - 8, diagExp)
                if np.sum(~np.isnan(corner)) > 100:
                    cornerpcc = pcc(corner[~np.isnan(corner)], looptpl[~np.isnan(corner)])[0]
                else:
                    cornerpcc = -1
                if cornerpcc>0.3:
                    result.append((l,r))
                else:
                    badnum+=1
    # print('add back',len(result))
    # print('bad',badnum)
    return result

def summary(HqTads,LqTads,data,chr,resol,minRatio=1):
    nestedScore=[]
    score = []
    Hq = []
    chrs = []
    startbp=[]
    endbp=[]
    level=[]
    # compute nestedScore , i.e. TADScore without masking subTADs
    for l,r in HqTads:
        start = max([0, 2 * l - r - 1])
        end = 2 * r - l + 1
        _nestedScore = Delta(data=data[start:end, start:end].copy(), offset=start, left=l, right=r,minRatio=minRatio, mask=None)
        nestedScore.append(_nestedScore)
        chrs.append(chr)
        Hq.append(1)
        startbp.append(l*resol)
        endbp.append(r*resol)
    for l, r in LqTads:
        start = max([0, 2 * l - r - 1])
        end = 2 * r - l + 1
        _nestedScore = Delta(data=data[start:end, start:end].copy(), offset=start, left=l, right=r, minRatio=minRatio,mask=None)
        nestedScore.append(_nestedScore)
        chrs.append(chr)
        Hq.append(0)
        startbp.append(l * resol)
        endbp.append(r * resol)
    # compute TADScore, i.e. TADScore by masking subTADs
    for tad in HqTads:
        nestTADs=[]
        for tad2 in HqTads:
            if tad2[0]>=tad[0] and tad2[1]<=tad[1]:
                if tad2[0]==tad[0] and tad2[1] == tad[1]:
                    pass
                else:
                    nestTADs.append(tad2)
        l, r = tad[0],tad[1]
        start = np.max([0, 2 * l - r - 1])
        end = 2 * r - l + 1
        _score = Delta(data=data[start:end, start:end].copy(), offset=start, left=l, right=r, minRatio=minRatio,mask=nestTADs)
        score.append(_score)
    for tad in LqTads:
        nestTADs = []
        for tad2 in HqTads:
            if tad2[0] >= tad[0] and tad2[1] <= tad[1]:
                if tad2[0] == tad[0] and tad2[1] == tad[1]:
                    pass
                else:
                    nestTADs.append(tad2)
        l, r = tad[0], tad[1]
        start = np.max([0, 2 * l - r - 1])
        end = 2 * r - l + 1
        _score = Delta(data=data[start:end, start:end].copy(), offset=start, left=l, right=r, minRatio=minRatio,
                       mask=nestTADs)
        score.append(_score)
    # compute TAD level
    levelPerBp = np.zeros(data.shape[0])-1
    tad2level={}
    tads = HqTads.copy()
    while len(tads) > 0:
        tad, tads = popOneL0TAD(tads)
        levelPerBp[tad[0]:tad[1]+1]+=1
        tad2level[(tad[0],tad[1])]=np.max(levelPerBp[tad[0]:tad[1]+1])
    tads = LqTads.copy()
    while len(tads) > 0:
        tad, tads = popOneL0TAD(tads)
        levelPerBp[tad[0]:tad[1] + 1] += 1
        tad2level[(tad[0], tad[1])] = np.max(levelPerBp[tad[0]:tad[1] + 1])
    for tad in HqTads:
        level.append(tad2level[(tad[0],tad[1])])
    for tad in LqTads:
        level.append(tad2level[(tad[0], tad[1])])
    return pd.DataFrame.from_dict({'chrom':chrs,'startbp':startbp,'endbp':endbp,'score':score,'nestedScore':nestedScore,'Hq':Hq,'level':level})

eps = np.finfo(float).eps

cached={}
counts={'hitcache':0,'misscache':0}

def cacheDelta(func):
    @functools.wraps(func)
    def lazzyDelta(**kwargs):
        tad = (kwargs['left'],kwargs['right'])
        if kwargs['mask'] is not None:
            mask = tuple(dict.fromkeys(kwargs['mask']))
        else:
            mask=()
        if tad in cached and mask in cached[tad]:
            counts['hitcache']+=1
            return cached[tad][mask]
        val=func(**kwargs)
        if tad not in cached:
            counts['misscache']+=1
            cached[tad]={}
        cached[tad][mask]=val
        return val

    return lazzyDelta

@jit(nopython=True, parallel=True,error_model='python')
def DeltaNB(diags, left, right,w, minRatio=1.1):

    scores = np.zeros(w)*np.nan
    N = 0

    for diag in numba.prange(1, w):
        if w>50 and diag%5!=1:
            continue
        crossIF1 = diags[max([0,left - diag]):left, diag]
        crossIF2 = diags[right -diag:right, diag]
        crossIF = np.concatenate((crossIF1, crossIF2))
        crossIF=crossIF[~np.isnan(crossIF)] + eps
        withinIF = diags[left:right - diag, diag] + eps
        withinIF=withinIF[~np.isnan(withinIF)]

        ratio = np.outer(withinIF, 1 / crossIF)
        n1, n2 = ratio.shape
        win = np.count_nonzero(ratio > minRatio)
        loss = np.count_nonzero(ratio < 1 / minRatio)
        # loss = np.count_nonzero(ratio <  minRatio)

        # n1 = len(withinIF)
        # n2 = len(crossIF)
        # if n1>0 and n2>0:
        #     win=timesOfAgeB(withinIF,crossIF)
        #     loss = timesOfAgeB(crossIF,withinIF)
        # else:
        #     win=1
        #     loss=1

        score = win - loss
        scale = n1*n2
        # scale=win + loss + eps
        score = score / scale * (n1 + n2)
        N += (n1 + n2)
        scores[diag]= score
    return scores[1:],N

@cacheDelta
def Delta(data, offset, left, right, minRatio=1.1, mask=None):
    data = data.copy()

    s = data.shape[0]
    if mask is not None:
        for i in range(len(mask)):
            l,r = mask[i]
            data[np.max(np.asarray([0, l - offset])):r - offset + 1, np.max(np.asarray([0, l - offset])):r - offset + 1] = np.nan
            data[np.max(np.asarray([0, l - offset - 4])):l - offset + 5,
            np.max(np.asarray([0, r - offset - 4])):r - offset + 5] = np.nan  # mask dot corner

    left = left - offset
    right = right - offset
    w = right - left
    diags = np.zeros((s, w + 2)) * np.nan

    for j in range(np.min(np.asarray([w + 2, s]))):
        diagj = np.diagonal(data, j)
        diags[:len(diagj), j] = diagj
    scores,N=DeltaNB(diags, left, right, w, minRatio)
    return np.nansum(np.asarray(scores)) / (N+eps)

def merge(data, TADs, distance=50000, minRatio=1.1):
    TADs = np.asarray(TADs)
    mergedTADs = []
    posTree = KDTree(TADs, leaf_size=30, metric='chebyshev')
    NNindexes, NNdists = posTree.query_radius(TADs, r=distance, return_distance=True)

    for i in range(len(NNindexes)):
        if len(NNindexes[i]) > 1:
            bestScore = -np.inf
            bestIdx = -1
            for j in range(len(TADs[NNindexes[i]])):
                l, r = TADs[NNindexes[i]][j]
                offset = np.max([0, 2 * l - r - 1])
                end = 2 * r - l + 1
                s = Delta(data=data[offset:end, offset:end], offset=offset, left=l, right=r, minRatio=minRatio,mask=None)
                if s > bestScore:
                    bestScore = s
                    bestIdx = j
            mergedTADs.append(list(TADs[NNindexes[i][bestIdx]]))

        else:
            mergedTADs.append(list(TADs[NNindexes[i][0]]))
    _mergedTADs = []
    for l, r in mergedTADs:
        _mergedTADs.append((l, r))
    mergedTADs = list(set(_mergedTADs))
    if len(TADs) > len(mergedTADs):
        mergedTADs = merge(data, mergedTADs, minRatio)
    return mergedTADs

def dp(data, lefts, rights, minDelta=0.3, minRatio=1.1, resol=5000, maxTAD=3000000, minTAD=30000):
    maxTAD = maxTAD / resol
    minTAD = minTAD / resol
    # distance = distance / resol
    s = data.shape[0]

    boundaries = sorted(list(set(lefts + rights)))
    lefts = set(lefts)
    rights = set(rights)
    if boundaries[0] not in lefts:
        lefts.add(0)
        boundaries.insert(0, 0)
    if boundaries[-1] not in rights:
        rights.add(s - 1)
        boundaries.append(s - 1)
    n = len(boundaries)
    S = np.zeros((n, n))
    pairS = np.zeros((n, n))
    T = {}
    L = {}

    for i in range(len(boundaries)):
        if boundaries[i] not in T:
            T[boundaries[i]] = {}
            L[boundaries[i]] = {}
    K = np.zeros((n, n), dtype=int) - 1
    # Initialization
    for i in range(n - 1):
        if boundaries[i] not in T:
            T[boundaries[i]] = {}
            L[boundaries[i]] = {}
        if boundaries[i + 1] not in T[boundaries[i]]:
            T[boundaries[i]][boundaries[i + 1]] = []
            L[boundaries[i]][boundaries[i + 1]] = 0
        if boundaries[i] in lefts and boundaries[i + 1] in rights and maxTAD > boundaries[i + 1] - boundaries[
            i] > minTAD:
            offset = np.max([0, 2 * boundaries[i] - boundaries[i + 1] - 1])
            end = 2 * boundaries[i + 1] - boundaries[i] + 1
            s = Delta(data=data[offset:end, offset:end], offset=offset, left=boundaries[i], right=boundaries[i + 1], mask=None,minRatio=minRatio)
            if s > minDelta:
                S[i, i + 1] = s
                pairS[i, i + 1] = s
                T[boundaries[i]][boundaries[i + 1]].append((boundaries[i], boundaries[i + 1]))
                L[boundaries[i]][boundaries[i + 1]] = 1

    # Foward pass
    for diag in tqdm(range(2, n)):
        for i in range(n):
            j = i + diag
            if j >= n:
                break

            bestScore = -np.inf
            bestTAD = []
            bestK = -1
            bestPairS = -1
            _pairS = 0
            _level = 0
            bestLevel  = 0

            for k in range(i + 1, j):
                s = S[i, k] + S[k, j]
                nestTad = T[boundaries[i]][boundaries[k]] + T[boundaries[k]][boundaries[j]]
                _level = np.max([L[boundaries[i]][boundaries[k]], L[boundaries[k]][boundaries[j]]])

                if boundaries[i] in lefts and boundaries[j] in rights and maxTAD > boundaries[j] - boundaries[
                            i] > minTAD:
                    offset = np.max([0, 2 * boundaries[i] - boundaries[j] - 1])
                    end = 2 * boundaries[j] - boundaries[i] + 1

                    s_delta = Delta(data=data[offset:end, offset:end], offset=offset, left=boundaries[i], right=boundaries[j],
                                    minRatio=minRatio, mask=nestTad)
                    if s_delta > minDelta:
                        largestMetaTADL = -np.inf
                        for _l, _r in nestTad:
                            if _r - _l > largestMetaTADL:
                                largestMetaTADL = _r - _l
                        if largestMetaTADL / (boundaries[j] - boundaries[i]) < 10:#0.9:
                            nestTad = [(boundaries[i], boundaries[j])]
                            s = s + np.nanmax([s_delta, 0])
                            _pairS = np.nanmax([s_delta, 0])
                            _level += 1

                if s > bestScore:
                    bestScore = s
                    bestTAD = nestTad
                    bestK = k
                    bestLevel = _level
                    bestPairS = _pairS
            S[i, j] = bestScore
            K[i, j] = bestK
            pairS[i, j] = bestPairS
            T[boundaries[i]][boundaries[j]] = bestTAD
            L[boundaries[i]][boundaries[j]] = bestLevel

    # Backtracking
    finalTADs = []
    TADlevels = {}
    unvisted = [[0, n - 1]]
    while len(unvisted) > 0:
        i, j = unvisted[0]
        unvisted.remove([i, j])
        if len(T[boundaries[i]][boundaries[j]]) == 1 and T[boundaries[i]][boundaries[j]][0][0] == boundaries[i] and T[boundaries[i]][boundaries[j]][0][1] == boundaries[j]:
            TADlevels[(boundaries[i], boundaries[j])] = L[boundaries[i]][boundaries[j]]
        finalTADs = finalTADs + T[boundaries[i]][boundaries[j]]
        k = K[i, j]
        if k != -1:
            unvisted.append([i, k])
            unvisted.append([k, j])

    finalTADs = list(set(finalTADs))

    # addTADs = addback(lefts,rights, data,minTAD,maxTAD)
    finalTADs, badTads = cleanTAD(finalTADs, data)
    # for l, r in addTADs:
    #     if (l,r) not in finalTADs:
    #         finalTADs.append((l, r))
    finalLefts = []
    finalRights = []
    unused = {'left': [], 'right': []}
    for l, r in finalTADs:
        finalLefts.append(l)
        finalRights.append(r)
    finalLefts = set(finalLefts)
    finalRights = set(finalRights)
    unused['left'] = lefts - finalLefts
    unused['right'] = rights - finalRights
    print(unused)
    # badTads=[]
    return finalTADs,badTads

    # finalTADs = merge(data, finalTADs, distance, minRatio)
    scores = []
    levels = []
    for l, r in finalTADs:
        scores.append(pairS[boundaries.index(l), boundaries.index(r)])
        if (l,r) in TADlevels:
            levels.append(TADlevels[(l,r)])
        else:
            levels.append(-1)

    return finalTADs, S[0, n - 1], scores,levels

@click.command()
@click.option('--resol', default=5000,type=int, help='resolution [5000]')
@click.option('--mintad', default=50000, type=int, help='min TAD size [50000]')
@click.option('--maxtad', default=3000000, type=int, help='max TAD size [3000000]')
# @click.option('--distance', default=50000, type=int, help='max distance for merging two TADs [50000]')
@click.option('--ratio', default=1.2, type=float, help='min ratio for comparision [1.2]')
@click.option('--delta', default=0.1, type=float, help='min score for forming a possible pair [0.1]')
# @click.option('--alpha', default=0.05, type=float, help='alpha-significant in FDR [0.05]')
@click.option('--chr', default=None, help='comma separated chromosomes')
@click.argument('coolfile', type=str, default=None, required=True)
@click.argument('boundaryfile', type=str, default=None, required=True)
@click.argument('prefix', type=str, default=None, required=True)
def assembly(delta, ratio, resol, maxtad, mintad, boundaryfile, coolfile, prefix, chr):
    '''
    Pairing left and right boundaries into nested TADs.
    '''
    if cooler.fileops.is_cooler(coolfile):
        c = cooler.Cooler(coolfile)
        if c.info['bin-size']!=resol:
            print('contact map at '+str(resol)+' does not exist!\nGood bye!')
            sys.exit(0)
    else:
        try:
            c = cooler.Cooler(coolfile + '::/resolutions/' + str(resol))
        except:
            print('contact map at '+str(resol)+' does not exist!\nGood bye!')
            sys.exit(0)

    print("analysis at ",resol, 'resolution')
    boundaryfile = pd.read_csv(boundaryfile, sep='\t', header=None)
    boundaryfile[[1,2]]//=resol
    results = {'chrom':[],'start':[],'end':[],'score':[],'level':[]}
    if chr is None:
        chr = list(set(boundaryfile[0]))
    else:
        chr = chr.split(',')

    pdresults=[]
    for _chr in chr:
        lefts = list(boundaryfile[(boundaryfile[0]==_chr)&(boundaryfile[4]==1)][1].to_numpy())
        rights = list(boundaryfile[(boundaryfile[0] == _chr) & (boundaryfile[6] == 1)][1].to_numpy())
        mat = c.matrix(balance=True,sparse=False).fetch(_chr)#.tocsr()

        # y=500
        # mat=mat[:y,:y]
        # lefts =np.asarray(lefts)
        # rights =np.asarray(rights)
        #
        # lefts = list(lefts[lefts < y])
        # rights = list(rights[rights < y])

        # _TADs, _score, _scores,levels = dp(mat, lefts, rights, delta, ratio, resol, maxtad, mintad, distance)
        hqTAD,lqTAD = dp(mat, lefts, rights, delta, ratio, resol, maxtad, mintad)
        del mat

        mat = c.matrix(balance=True, sparse=False).fetch(_chr)  # .tocsr()

        pdresults.append(summary(hqTAD, lqTAD, mat, _chr, resol, minRatio=ratio))

    result=pd.concat(pdresults)
    result[(result['Hq']==1) & (result['score']>0)].to_csv(prefix+'_hq.bed',sep='\t',index=False,header=False)
    result.to_csv(prefix+'_all.bed', sep='\t', index=False, header=True)


from decouple import config
from deta import Deta
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

DETA_PROJECT_KEY = config('DETA_PROJECT_KEY')
DB_NAME = config('DB_NAME')

app = FastAPI()
deta = Deta(DETA_PROJECT_KEY)
db = deta.Base(DB_NAME)

class Driver(BaseModel):
    key: str
    link: str

@app.get("/wacom-drivers", status_code=200)
def get_drivers():
    response = db.fetch()
    drivers = response.items
    return drivers

@app.get("/wacom-drivers/{key}", status_code=200)
def get_driver(key: str):
    driver = db.get(key)
    if driver is None:
        raise HTTPException(status_code=404, detail="Driver not found")
    return driver


import json

from websocket import create_connection

ws = create_connection("wss://api.bitfinex.com/ws/2")
#ws.connect("wss://api2.bitfinex.com:3000/ws")
ws.send(json.dumps({
    "event": "subscribe",
    "channel": "book",
    "symbol": "tBTCUSD",
}))

while True:
    result = ws.recv()
    result = json.loads(result)
    xxx = result
    print(xxx)
    #print(result[1])

ws.close()


#!/usr/bin/env python3
import argparse
from botocore.exceptions import ClientError
import datetime

module_info = {
    # Name of the module (should be the same as the filename)
    'name': 'aws__enum_spend',

    # Name and any other notes about the author
    'author': 'Chris Farris <chris@room17.com>',

    # Category of the module. Make sure the name matches an existing category.
    'category': 'ENUM',

    # One liner description of the module functionality. This shows up when a user searches for modules.
    'one_liner': 'Enumerates account spend by service.',

    # Full description about what the module does and how it works
    'description': 'Display what services the account uses and how much is spent. Data is pulled from CloudWatch metrics and the AWS/Billing Namespace.',

    # A list of AWS services that the module utilizes during its execution
    'services': ['IAM'],

    # For prerequisite modules, try and see if any existing modules return the data that is required for your module before writing that code yourself, that way, session data can stay separated and modular.
    'prerequisite_modules': [],

    # External resources that the module depends on. Valid options are either a GitHub URL (must end in .git) or single file URL.
    'external_dependencies': [],

    # Module arguments to autocomplete when the user hits tab
    'arguments_to_autocomplete': [],
}

parser = argparse.ArgumentParser(add_help=False, description=module_info['description'])

# Main is the first function that is called when this module is executed
def main(args, pacu_main):
    session = pacu_main.get_active_session()

    ###### Don't modify these. They can be removed if you are not using the function.
    args = parser.parse_args(args)
    print = pacu_main.print
    ######

    # All the billing seems to be in us-east-1. YMMV
    cwm_client = pacu_main.get_boto3_client('cloudwatch', "us-east-1")

    services = []
    service_spend = {}

    try:
        response = cwm_client.list_metrics(
            Namespace='AWS/Billing',
            MetricName='EstimatedCharges'
        )
        metrics = response['Metrics']
        for m in metrics:
            for d in m['Dimensions']:
                if d['Name'] == "ServiceName":
                    services.append(d['Value'])
        if len(services) == 0:
            print('\nNo services found. Unable to determine account spend.\n')
            return
    except ClientError as e:
        print("ClientError getting spend: {}".format(e))
        return({"error": "<unauthorized>"})

    for s in services:
        try:
            print("Retrieving metrics for service {}...".format(s))
            response = cwm_client.get_metric_statistics(
                Namespace='AWS/Billing',
                MetricName='EstimatedCharges',
                Dimensions=[
                    {
                        'Name': 'ServiceName',
                        'Value': s
                    },
                    {
                        'Name': 'Currency',
                        'Value': 'USD'
                    },
                ],
                StartTime=datetime.datetime.now() - datetime.timedelta(hours=6),
                EndTime=datetime.datetime.now(),
                Period=21600,  # 6 hours
                Statistics=['Maximum'],
                Unit='None'
            )
            if len(response['Datapoints']) == 0:
                service_spend[s] = 0
            else:
                service_spend[s] = response['Datapoints'][0]['Maximum']
        except KeyError as e:
            print("KeyError getting spend: {} -- Response: {}".format(e, response))
        except IndexError as e:
            print("IndexError getting spend: {} -- Response: {}".format(e, response))
        except ClientError as e:
            print("ClientError getting spend: {}".format(e))

    session.update(pacu_main.database, AccountSpend=service_spend)

    return service_spend

def summary(data, pacu_main):
    out = "Account Spend:\n"
    for key in sorted(data.keys(), key=lambda x: data[x], reverse=True):
        out += "        {:<30}: {:>10.2f} (USD)\n".format(key, data[key])
    return out


import ticketsys

while True:
    print("\n\n\n")
    print("Selection Menu\n")
    print("1:Create a ticket ")
    print("2:submit")
    print("3:Provide a response")
    print("4:reopen a ticket")
    print("5:display stats")
    print("6:quit")
    smenu=input()
    if smenu==1:
        ""
    elif smenu==2:
        ""
    elif smenu==3:
        ""
    elif smenu==4:
        ""
    elif smenu==5:
        ""
    elif smenu==6:
        ""
    else:
        print("NOT AN OPTION")


from tqdm import tqdm
from random import randint

def insertion_sort(array):
	for i in tqdm(range(1, len(array))):
		current = array[i]
		j = i - 1
		while j >= 0 and array[j] > current:
			array[j + 1] = array[j]
			j -= 1
		array[j + 1] = current
	return array

nums = [randint(1,20000) for _ in range(20000)]
print(insertion_sort(nums))


import time
import serial
import serial.tools.list_ports
import threading

def send_servo(digree):
    sendData = f"SERVO={digree}\n"
    my_serial.write( sendData.encode() )

def send_buzzer(freq):
    sendData = f"BUZZER={freq}\n"
    my_serial.write( sendData.encode() )

def send_fnd(data):
    sendData = f"FND={data}\n"
    my_serial.write( sendData.encode() )

def serial_read_thread():
    while True:
        read_data = my_serial.readline()
        print(read_data.decode())

def main():
    try:
        while True:
            send_servo(0)
            send_buzzer(261)
            send_fnd(1234)
            time.sleep(1.0)

            send_servo(90)
            send_buzzer(293)
            send_fnd(5678)
            time.sleep(1.0)

            send_servo(180)
            send_buzzer(329)
            send_fnd(3.14)
            time.sleep(1.0)

            send_servo(90)
            send_buzzer(349)
            send_fnd(12.34)
            time.sleep(1.0)
    except KeyboardInterrupt:
        send_buzzer(0)

if __name__ == '__main__':

    ports = list(serial.tools.list_ports.comports())
    for p in ports:
        if 'Arduino Uno' in p.description:
            print(f"{p} 포트에 연결하였습니다.")
            my_serial = serial.Serial(p.device, baudrate=9600, timeout=1.0)
            time.sleep(2.0)

    t1 = threading.Thread(target=serial_read_thread)
    t1.daemon = True
    t1.start()

    main()

    my_serial.close()

한끝아파52 소스코드/02.파이썬코드/7_2.파이썬으로 아두이노 서보모터,피에조부저,FND 제어하기/main7_2_2.py


class VolConst():
    ph_ashf = 'ASH'
    ph_lava = 'LAVA'
    ph_laha = 'LAHAR'
    ph_pyro = 'PYRO'

    ty_text = 'text'
    ty_twkt = 'text-to-wkt'
    ty_open = 'openquake'
    ty_shap = 'shape'
    ty_swkt = 'shape-to-wkt'

    g = 9.80665

    epsg_out = '4326'


import math
def calculate_volume(diameter, height):
    # the formula for volume of cone V = πr^2h/3
    radius = diameter/2
    volume = math.pi * radius ** 2 * (height / 3)
    print("Mr Survivor: the volume of a cone with a diameter of {} and height of {} is {} herfs ^3".format(diameter, height, volume))

def calculate_surfacearea(diameter, height):
     #The formula for surface area A = πr(r + sqr(h^2 + r^2))
     #The base area A = πr^2

     radius = diameter/2
     base_area = math.pi * radius ** 2
     surface_area = math.pi * radius * (radius + math.sqrt(height ** 2 + radius ** 2))
     surface_area_side = surface_area - base_area
     print("The surface area = {} herfs ^2 including the base.".format(surface_area))
     print("The surface area of just the sides = {} herfs ^2 ".format(surface_area_side))

def main():
    #prompt the user to enter non negative diameter and height
    diameter = float(input("Input base Diameter (herfs): "))
    #check that the diameter is non negative
    while diameter < 0:
        diameter = float(input("Input base Diameter (herfs): "))

    #check that the height is non negative
    height = float(input("Input height (herfs): "))
    while height < 0 :
        height = float(input("Input height (herfs): "))

    #call function to calculate volume and area respectively and passing diameter and height to them
    calculate_volume(diameter, height)
    calculate_surfacearea(diameter, height)
#call the main function
main()


import pytest
from django.test import TestCase
from django.urls import reverse

from rard.research.models import (
    Antiquarian,
    Book,
    Fragment,
    Testimonium,
    TextObjectField,
    Work,
)
from rard.research.models.base import FragmentLink, TestimoniumLink
from rard.research.models.work import collate_unknown

pytestmark = pytest.mark.django_db

class TestWork(TestCase):
    def test_creation(self):
        # can create with a name only
        data = {
            "name": "Work Name",
            "subtitle": "Subtitle",
        }
        work = Work.objects.create(**data)
        for key, val in data.items():
            self.assertEqual(getattr(work, key), val)

    def test_ordering(self):
        # ordering should be wrt to the authors of the work
        anta = Antiquarian.objects.create(name="AAA", re_code="aaa")
        antb = Antiquarian.objects.create(name="BBB", re_code="bbb")
        antc = Antiquarian.objects.create(name="CCC", re_code="ccc")

        worka = Work.objects.create(name="aaa")
        workb = Work.objects.create(name="bbb")
        workc = Work.objects.create(name="ccc")

        # 1. check anonymous ordering
        self.assertEqual(
            [w for w in Work.objects.exclude(unknown=True)], [worka, workb, workc]
        )

        antc.works.add(workb)
        # this should now be last with anon works at the start
        self.assertEqual(
            [w for w in Work.objects.exclude(unknown=True)], [worka, workc, workb]
        )

        # the name of the antiquarian should put work c second
        antb.works.add(workc)
        self.assertEqual(
            [w for w in Work.objects.exclude(unknown=True)], [worka, workc, workb]
        )
        # even if we also add antc as an author of workc, as the name of
        # antiquarian antb should govern the order
        antc.works.add(workb)
        self.assertEqual(
            [w for w in Work.objects.exclude(unknown=True)], [worka, workc, workb]
        )

        # now, put worka as a work of anta and this should be first in the
        # list as the name of the antiquarian and then the work should be
        # ahead of the others
        anta.works.add(worka)
        self.assertEqual(
            [w for w in Work.objects.exclude(unknown=True)], [worka, workc, workb]
        )

        # final test - antiquarian name takes precedence
        anta.works.set([workc])
        antb.works.set([workb])
        antc.works.set([worka])
        self.assertEqual(
            [w for w in Work.objects.exclude(unknown=True)], [workc, workb, worka]
        )

    def test_required_fields(self):
        self.assertFalse(Work._meta.get_field("name").blank)
        self.assertTrue(Work._meta.get_field("subtitle").blank)

    def test_display_anonymous(self):
        # the __str__ function should show the name
        data = {
            "name": "Work Name",
            "subtitle": "Subtitle",
        }
        work = Work.objects.create(**data)
        self.assertEqual(str(work), "Anonymous: %s" % data["name"])

    def test_display(self):
        # the __str__ function should show the name
        data = {
            "name": "Work Name",
            "subtitle": "Subtitle",
        }
        work = Work.objects.create(**data)
        antiquarian1 = Antiquarian.objects.create(
            name="John Smith", re_code="smitre001"
        )
        antiquarian1.works.add(work)
        self.assertEqual(str(work), "%s: %s" % (antiquarian1.name, data["name"]))
        # now test mutiple antiquarians
        antiquarian2 = Antiquarian.objects.create(
            name="Joe Bloggs", re_code="blogre002"
        )
        antiquarian2.works.add(work)
        self.assertEqual(
            str(work),
            "%s, %s: %s"
            % (
                antiquarian2.name,
                antiquarian1.name,
                data["name"],
            ),
        )

    def test_work_can_belong_to_multiple_antiquarians(self):
        # works can belong to multiple antiquarians
        data = {
            "name": "Work Name",
            "subtitle": "Subtitle",
        }
        work = Work.objects.create(**data)
        length = 10
        for counter in range(0, length):
            antiquarian = Antiquarian.objects.create(
                name="John Smith", re_code="smitre%03d" % counter
            )
            antiquarian.works.add(work)

        self.assertEqual(work.antiquarian_set.count(), length)

    # def test_disallow_reverse_relator_addition(self):
    #     # to maintain links between antiquarians and fragments when
    #     # works are added and removed, we only support addition of
    #     # works to antiquarians in one direction
    #     antiquarian = Antiquarian.objects.create(
    #         name='John Smith', re_code='smitre001'
    #     )
    #     data = {
    #         'name': 'Work Name',
    #         'subtitle': 'Subtitle',
    #     }
    #     work = Work.objects.create(**data)

    #     # allowed
    #     antiquarian.works.add(work)

    #     antiquarian.works.clear()

    #     # not allowed
    #     with self.assertRaises(IntegrityError):
    #         work.antiquarian_set.add(antiquarian)

    def test_deleting_work_leaves_antiquarian(self):
        # check we are using uuids as primary keys
        data = {
            "name": "Work Name",
            "subtitle": "Subtitle",
        }
        work = Work.objects.create(**data)
        antiquarian = Antiquarian.objects.create(name="John Smith")
        antiquarian.works.add(work)
        antiquarian_pk = antiquarian.pk
        # delete the work
        work.delete()
        # we should still have the antiquarian
        Antiquarian.objects.get(pk=antiquarian_pk)

    def test_order_by_name(self):
        names = [
            "Self Hypnosis for beginners",
            "Fly Fishing",
            "Django development",
            "Creative Writing",
        ]
        for counter, name in enumerate(names):
            Work.objects.create(name=name)

        db_names = []
        for a in Work.objects.all():
            db_names.append(a.name)

        self.assertEqual(db_names, sorted(names))

    def test_get_absolute_url(self):
        data = {
            "name": "Work Name",
            "subtitle": "Subtitle",
        }
        work = Work.objects.create(**data)
        self.assertEqual(
            work.get_absolute_url(), reverse("work:detail", kwargs={"pk": work.pk})
        )

    def test_fragment_methods(self):
        data = {
            "name": "Work Name1",
            "subtitle": "Subtitle",
        }
        work1 = Work.objects.create(**data)
        data = {
            "name": "Work Name2",
            "subtitle": "Subtitle",
        }
        work2 = Work.objects.create(**data)

        # create some fragments
        for i in range(0, 10):
            data = {
                "name": "name{}".format(i),
            }
            Fragment.objects.create(**data)

        # link to work1
        for fragment in Fragment.objects.all():
            FragmentLink.objects.create(
                work=work1,
                fragment=fragment,
                definite_work=True,
            )
        # shoud appear in work1's definite fragments only
        self.assertEqual(
            [x.pk for x in work1.definite_fragments()],
            [x.pk for x in Fragment.objects.all()],
        )
        self.assertEqual(0, len(work1.possible_fragments()))
        # and not in work2's
        self.assertEqual(0, len(work2.definite_fragments()))
        self.assertEqual(0, len(work2.possible_fragments()))

        # make these links possible...
        FragmentLink.objects.update(definite_work=False)
        # shoud appear in work1's possible fragments only
        self.assertEqual(
            [x.pk for x in work1.possible_fragments()],
            [x.pk for x in Fragment.objects.all()],
        )
        self.assertEqual(0, len(work1.definite_fragments()))
        # and not in work2's
        self.assertEqual(0, len(work2.definite_fragments()))
        self.assertEqual(0, len(work2.possible_fragments()))

        # switch all fragment links to work2 and make definite...
        FragmentLink.objects.update(work=work2, definite_work=True)

        # shoud appear in work2's definite fragments only
        self.assertEqual(
            [x.pk for x in work2.definite_fragments()],
            [x.pk for x in Fragment.objects.all()],
        )
        self.assertEqual(0, len(work2.possible_fragments()))
        # and not in work1's
        self.assertEqual(0, len(work1.definite_fragments()))
        self.assertEqual(0, len(work1.possible_fragments()))

        # finally make possible and check...
        FragmentLink.objects.update(definite_work=False)
        # shoud appear in work2's possible fragments only
        self.assertEqual(
            [x.pk for x in work2.possible_fragments()],
            [x.pk for x in Fragment.objects.all()],
        )
        self.assertEqual(0, len(work2.definite_fragments()))
        # and not in work1's
        self.assertEqual(0, len(work1.definite_fragments()))
        self.assertEqual(0, len(work1.possible_fragments()))

    def test_testimonium_methods(self):
        data = {
            "name": "Work Name1",
            "subtitle": "Subtitle",
        }
        work1 = Work.objects.create(**data)
        data = {
            "name": "Work Name2",
            "subtitle": "Subtitle",
        }
        work2 = Work.objects.create(**data)

        # create some testimonia
        for i in range(0, 10):
            data = {
                "name": "name{}".format(i),
            }
            Testimonium.objects.create(**data)

        # link to work1
        for testimonium in Testimonium.objects.all():
            TestimoniumLink.objects.create(
                work=work1,
                testimonium=testimonium,
                definite_work=True,
            )
        # shoud appear in work1's definite testimonia only
        self.assertEqual(
            [x.pk for x in work1.definite_testimonia()],
            [x.pk for x in Testimonium.objects.all()],
        )
        self.assertEqual(0, len(work1.possible_testimonia()))
        # and not in work2's
        self.assertEqual(0, len(work2.definite_testimonia()))
        self.assertEqual(0, len(work2.possible_testimonia()))

        # make these links possible...
        TestimoniumLink.objects.update(definite_work=False)
        # shoud appear in work1's possible testimonia only
        self.assertEqual(
            [x.pk for x in work1.possible_testimonia()],
            [x.pk for x in Testimonium.objects.all()],
        )
        self.assertEqual(0, len(work1.definite_testimonia()))
        # and not in work2's
        self.assertEqual(0, len(work2.definite_testimonia()))
        self.assertEqual(0, len(work2.possible_testimonia()))

        # switch all testimonium links to work2 and make definite...
        TestimoniumLink.objects.update(work=work2, definite_work=True)

        # shoud appear in work2's definite testimonia only
        self.assertEqual(
            [x.pk for x in work2.definite_testimonia()],
            [x.pk for x in Testimonium.objects.all()],
        )
        self.assertEqual(0, len(work2.possible_testimonia()))
        # and not in work1's
        self.assertEqual(0, len(work1.definite_testimonia()))
        self.assertEqual(0, len(work1.possible_testimonia()))

        # finally make possible and check...
        TestimoniumLink.objects.update(definite_work=False)
        # shoud appear in work2's possible testimonia only
        self.assertEqual(
            [x.pk for x in work2.possible_testimonia()],
            [x.pk for x in Testimonium.objects.all()],
        )
        self.assertEqual(0, len(work2.definite_testimonia()))
        # and not in work1's
        self.assertEqual(0, len(work1.definite_testimonia()))
        self.assertEqual(0, len(work1.possible_testimonia()))

    def test_introduction_created_with_work(self):
        data = {
            "name": "workname",
            "subtitle": "Subtitle",
        }
        work = Work.objects.create(**data)
        self.assertIsNotNone(work.introduction.pk)
        self.assertEqual(
            TextObjectField.objects.get(pk=work.introduction.pk), work.introduction
        )

    def test_introduction_deleted_with_book(self):
        data = {
            "name": "workname",
            "subtitle": "Subtitle",
        }
        work = Work.objects.create(**data)
        introduction_pk = work.introduction.pk
        work.delete()
        with self.assertRaises(TextObjectField.DoesNotExist):
            TextObjectField.objects.get(pk=introduction_pk)

    def test_collate_unknown(self):
        data = {
            "name": "workname",
            "subtitle": "Subtitle",
        }
        w = Work.objects.create(**data)
        original_unknown = w.unknown_book
        additional_unknown = Book.objects.create(
            unknown=True, subtitle="Unknown Book", work=w
        )
        additional_unknown.save()
        self.assertEqual(w.book_set.filter(unknown=True).count(), 2)
        collate_unknown(w)
        self.assertEqual(w.book_set.filter(unknown=True).count(), 1)
        self.assertEqual(original_unknown.pk, w.unknown_book.pk)

class TestBook(TestCase):
    def setUp(self):
        self.work = Work.objects.create(name="book_name")

    def test_creation(self):
        # test happy path
        data = {"number": "1", "subtitle": "Subtitle", "work": self.work}
        book = Book.objects.create(**data)
        for key, val in data.items():
            self.assertEqual(getattr(book, key), val)

    def test_work_required(self):
        data_no_work = {
            "number": "1",
            "subtitle": "Subtitle",
        }
        with self.assertRaises(Book.work.RelatedObjectDoesNotExist):
            Book.objects.create(**data_no_work)

    def test_required_fields(self):
        self.assertTrue(Book._meta.get_field("number").blank)
        self.assertTrue(Book._meta.get_field("subtitle").blank)

    def test_get_absolute_url(self):
        # the work url is the absolute url for its book objects
        data = {
            "number": "1",
            "subtitle": "Subtitle",
            "work": self.work,
        }
        book = Book.objects.create(**data)
        self.assertEqual(
            book.get_absolute_url(), reverse("work:detail", kwargs={"pk": self.work.pk})
        )

    def test_display(self):
        # the __str__ function should show the name
        data = {
            "number": "1",
            "subtitle": "Subtitle",
            "work": self.work,
        }
        book = Book.objects.create(**data)
        self.assertEqual(str(book), "Book 1: Subtitle")
        book.number = ""
        self.assertEqual(str(book), "Subtitle")
        book.number = "1"
        book.subtitle = ""
        self.assertEqual(str(book), "Book 1")

    def test_introduction_created_with_book(self):
        data = {
            "number": "1",
            "subtitle": "Subtitle",
            "work": self.work,
        }
        book = Book.objects.create(**data)
        self.assertIsNotNone(book.introduction.pk)
        self.assertEqual(
            TextObjectField.objects.get(pk=book.introduction.pk), book.introduction
        )

    def test_introduction_deleted_with_book(self):
        data = {
            "number": "1",
            "subtitle": "Subtitle",
            "work": self.work,
        }
        book = Book.objects.create(**data)
        introduction_pk = book.introduction.pk
        book.delete()
        with self.assertRaises(TextObjectField.DoesNotExist):
            TextObjectField.objects.get(pk=introduction_pk)


from argparse import ArgumentError
import astropy.constants as ac
import astropy.units as au
import numpy as np
from scipy.optimize import brentq

# setting up some convenient conversion factors
_Gconst_cgs = ac.G.cgs.value
_pc_cgs = ac.pc.cgs.value
_kbol_cgs = ac.k_B.cgs.value

_surf_cgs = (ac.M_sun / ac.pc**2).cgs.value
_rho_cgs = (ac.M_sun / ac.pc**3).cgs.value
_sfr_cgs = (ac.M_sun / ac.kpc**2 / au.yr).cgs.value
_kms_cgs = (1 * au.km / au.s).cgs.value

_sigma_eff_models = dict()
_sigma_eff_models["tigress_mid"] = dict(sigma_0=9.8, expo=0.15, sigma_min=5)
_sigma_eff_models["tigress_avg"] = dict(sigma_0=12, expo=0.22, sigma_min=5)

_yield_models = dict()
_yield_models["tigress-classic"] = dict(Y0=10**3.86, expo=-0.212)
_yield_models["tigress-classic-decomp"] = dict(
    Yth0=10**4.45, expo_th=-0.506, Ytrb0=1.5 * 10**2.81, expo_trb=-0.06
)
_yield_models["tigress-ncr-decomp"] = dict(
    Yth0=10**4.7, expo_th=-0.5, Ytrb0=10**3.3, expo_trb=-0.1
)
_yield_models["tigress-ncr-decomp-Z01"] = dict(
    Yth0=10**5.2, expo_th=-0.5, Ytrb0=10**3.3, expo_trb=-0.1
)

def get_weight_gas(Sigma_gas):
    """weight due to gas self-gravity"""
    return 0.5 * np.pi * _Gconst_cgs * Sigma_gas**2

def get_weight_star(Sigma_gas, H_gas, Sigma_star, H_star):
    """weight due to stellar gravity"""
    return np.pi * _Gconst_cgs * Sigma_gas * Sigma_star * H_gas / (H_gas + H_star)

def get_weight_star_gaussian(Sigma_gas, H_gas, Sigma_star, H_star):
    """weight due to stellar gravity for Gaussian profiles"""
    return 2 * _Gconst_cgs * Sigma_gas * Sigma_star * np.arctan(H_gas / H_star)

def get_weight_dm(Sigma_gas, H_gas, Omega_d, zeta_d=1 / 3.0):
    """weight due to dark matter gravity"""
    return zeta_d * Sigma_gas * H_gas * Omega_d**2

def get_pressure(Sigma_gas, H_gas, sigma_eff):
    """total pressure"""
    return 0.5 * Sigma_gas / H_gas * sigma_eff**2

@np.vectorize
def get_weights(
    Sigma_gas, Sigma_star, Omega_d, H_star, sigma_eff, zeta_d=1 / 3.0, method="analytic"
):
    """calculate gas scale height and then each weight term

    Parameters
    ----------
    Sigma_gas: float or array-like
        gas surface density
    Sigma_star: float or array-like
        stellar surface density
    Omega_d: float or array-like
        galactic rotation speed
    H_star: float or array-like
        stellar scale height
    sigma_eff: str or float or array-like
        velocity dispersion
    zeta_d: float [1/3.]
        parameter for dm and gas distribution
    method: str [analytic or numerical]
        calculate scale height either analytically or numerically

    Returns
    -------
    H_gas: float or array-like
        gas scale height
    W_gas: float or array-like
        weight by gas
    W_star: float or array-like
        weight by star
    W_dm: float or array-like
        weight by dark matter

    Example
    -------
    >>> import prfm
    >>> H, W_gas, W_star, W_dm = prfm.get_weights(Sigma_gas, Sigma_star, Omega_d, H_star, sigma_eff)
    """
    H_gas = get_scale_height(
        Sigma_gas, Sigma_star, Omega_d, H_star, sigma_eff, zeta_d=zeta_d, method=method
    )

    W_gas = get_weight_gas(Sigma_gas)
    W_star = get_weight_star(Sigma_gas, H_gas, Sigma_star, H_star)
    W_dm = get_weight_dm(Sigma_gas, H_gas, Omega_d, zeta_d=zeta_d)

    return H_gas, W_gas, W_star, W_dm

def get_weight_star_thick(Sigma_gas, rho_star, sigma_eff):
    W_star = Sigma_gas * np.sqrt(2 * _Gconst_cgs * rho_star) * sigma_eff
    return W_star

def get_weights_thick(Sigma_gas, rho_star, sigma_eff):
    """calculate gas scale height and then each weight term for H_*>>H_gas

    Parameters
    ----------
    Sigma_gas: float or array-like
        gas surface density
    rho_star: float or array-like
        stellar volume density
    sigma_eff: str or float or array-like
        velocity dispersion

    Returns
    -------
    H_gas: float or array-like
        gas scale height
    W_gas: float or array-like
        weight by gas
    W_star: float or array-like
        weight by star
    W_dm: float or array-like
        weight by dark matter

    Note
    ----
    Used approximate formula Equation (7) in Ostriker and Kim (2022)
    """
    H_gas = get_scale_height_thick(Sigma_gas, rho_star, sigma_eff)

    W_gas = get_weight_gas(Sigma_gas)
    W_star = get_weight_star_thick(Sigma_gas, rho_star, sigma_eff)
    W_dm = 0.0 * H_gas / H_gas

    return H_gas, W_gas, W_star, W_dm

def get_weight_star_thin(Sigma_gas, Sigma_star, sigma_eff):
    W_star = np.pi * _Gconst_cgs * Sigma_gas * Sigma_star
    return W_star

def get_weights_thin(Sigma_gas, Sigma_star, sigma_eff):
    """calculate gas scale height and then each weight term for H_*>>H_gas

    Parameters
    ----------
    Sigma_gas: float or array-like
        gas surface density
    Sigma_star: float or array-like
        stellar surface density
    sigma_eff: str or float or array-like
        velocity dispersion

    Returns
    -------
    H_gas: float or array-like
        gas scale height
    W_gas: float or array-like
        weight by gas
    W_star: float or array-like
        weight by star
    W_dm: float or array-like
        weight by dark matter

    Note
    ----
    Used approximate formula Equation (7) in Ostriker and Kim (2022)
    """
    H_gas = get_scale_height_thin(Sigma_gas, Sigma_star, sigma_eff)

    W_gas = get_weight_gas(Sigma_gas)
    W_star = get_weight_star_thin(Sigma_gas, Sigma_star, sigma_eff)
    W_dm = 0.0 * H_gas / H_gas

    return H_gas, W_gas, W_star, W_dm

@np.vectorize
def get_weight_contribution(
    Sigma_gas, Sigma_star, Omega_d, H_star, sigma_eff, zeta_d=1 / 3.0, method="analytic"
):
    """Wrapper function to calculate the ratio of each weight term to the total weight

    Parameters
    ----------
    Sigma_gas: float or array-like
        gas surface density
    Sigma_star: float or array-like
        stellar surface density
    Omega_d: float or array-like
        galactic rotation speed
    H_star: float or array-like
        stellar scale height
    sigma_eff: str or float or array-like
        velocity dispersion
    zeta_d: float [1/3.]
        parameter for dm and gas distribution
    method: str [analytic or numerical]
        calculate scale height either analytically or numerically

    Returns
    -------
    f_gas: float or array-like
        weight by gas/total weight
    f_star: float or array-like
        weight by star/total weight
    f_dm: float or array-like
        weight by dark matter/total weight
    """
    H, wgas, wstar, wdm = get_weights(
        Sigma_gas, Sigma_star, Omega_d, H_star, sigma_eff, zeta_d=zeta_d, method=method
    )
    wtot = wgas + wstar + wdm

    return wgas / wtot, wstar / wtot, wdm / wtot

@np.vectorize
def get_scale_height_gas_only(*args, **kwargs):
    """analytic solution for gas only case"""
    Sigma_gas, Sigma_star, Omega_d, H_star, sigma_eff = args
    return sigma_eff**2 / (np.pi * _Gconst_cgs * Sigma_gas)

@np.vectorize
def get_scale_height_star_only(*args, **kwargs):
    """analytic solution for star only case"""
    Sigma_gas, Sigma_star, Omega_d, H_star, sigma_eff = args
    h = sigma_eff**2 / (4 * np.pi * _Gconst_cgs * Sigma_star)
    h_star = H_star / h
    return h * (1 + np.sqrt(1 + 2 * h_star))

@np.vectorize
def get_scale_height_dm_only(*args, **kwargs):
    """analytic solution for dark matter only case"""
    Sigma_gas, Sigma_star, Omega_d, H_star, sigma_eff = args
    zeta_d = kwargs["zeta_d"]
    return sigma_eff / np.sqrt(2 * zeta_d) / Omega_d

@np.vectorize
def get_scale_height_star_gas(*args, **kwargs):
    """analytic solution neglecting dark matter"""
    Sigma_gas, Sigma_star, Omega_d, H_star, sigma_eff = args
    H_gas_only = get_scale_height_gas_only(*args, **kwargs)
    eta_star = H_star / H_gas_only
    s_star = Sigma_star / Sigma_gas
    # h = ((1-eta_star)+np.sqrt((eta_star-1)**2+4*eta_star*(1+2.*s_star)))/(2*(1+2*s_star))
    h = (
        2
        * eta_star
        / ((eta_star - 1) + np.sqrt((eta_star + 1) ** 2 + 8 * s_star * eta_star))
    )
    return h * H_gas_only

@np.vectorize
def get_scale_height_star_gas_approx(*args, **kwargs):
    """analytic solution neglecting dark matter"""
    Sigma_gas, Sigma_star, Omega_d, H_star, sigma_eff = args
    H_gas_only = get_scale_height_gas_only(*args, **kwargs)
    eta_star = H_star / H_gas_only
    s_star = Sigma_star / Sigma_gas
    h = (1 + np.sqrt(1 + 8 * s_star * eta_star)) / (4 * s_star)
    return h * H_gas_only

@np.vectorize
def get_scale_height_dm_gas(*args, **kwargs):
    """analytic solution neglectic star"""
    H_gas_only = get_scale_height_gas_only(*args, **kwargs)
    H_dm_only = get_scale_height_dm_only(*args, **kwargs)
    eta_D = H_dm_only / H_gas_only
    h = (-(eta_D**2) + np.sqrt(eta_D**4 + 4 * eta_D**2)) * 0.5
    return h * H_gas_only

@np.vectorize
def get_sigma_eff(P, model="tigress_mid"):
    """pressure-dependent velocity dispersion model"""
    sigma_eff_model = _sigma_eff_models[model]
    sigma_0 = sigma_eff_model["sigma_0"]
    expo = sigma_eff_model["expo"]
    sigma_min = sigma_eff_model["sigma_min"]

    veld = sigma_0 * (1.0e-4 * P / _kbol_cgs) ** expo

    return np.clip(veld, sigma_min, None) * 1.0e5

@np.vectorize
def get_feedback_yield(P, model="tigress-classic"):
    """total feedback yield as a function of weight

    Paramerters
    -----------
    P : float
        midplane pressure or weight
    model : str
        tigress-classic for OK22
        tigress-NCR for K23

    Returns
    -------
    Ytot : float
        feedback yield in km/s
    """
    yield_model = _yield_models[model]
    if "Y0" in yield_model:
        Y0 = yield_model["Y0"]
        slope = yield_model["expo"]
        Ytot = Y0 * (P / _kbol_cgs) ** slope
    elif "Yth0" in yield_model:
        Yth0 = yield_model["Yth0"]
        expo_th = yield_model["expo_th"]
        Ytrb0 = yield_model["Ytrb0"]
        expo_trb = yield_model["expo_trb"]
        Ytot = Yth0 * (P / _kbol_cgs) ** expo_th + Ytrb0 * (P / _kbol_cgs) ** expo_trb
    return Ytot

@np.vectorize
def get_feedback_yield_comp(P, comp="th", model="tigress-classic-decomp"):
    """feedback yield of each component as a function of weight

    Paramerters
    -----------
    P : float
        midplane pressure or weight
    comp : str
        component name ['th','trb']
    model : str
        tigress-classic for OK22
        tigress-NCR for K23

    Returns
    -------
    Y : float
        feedback yield in km/s
    """
    yield_model = _yield_models[model]
    Y0 = yield_model["Y{}0".format(comp)]
    expo = yield_model["expo_{}".format(comp)]
    return Y0 * (P / _kbol_cgs) ** expo

def get_sfr(P, Ytot="tigress-classic"):
    """calculate SFR surface density using feedback yield

    Paramerters
    -----------
    P : float
        midplane pressure or weight
    Ytot : float, str
        constant value in km/s or
        P-dependent feedback yield model e.g., `tigress-classic`

    Returns
    -------
    sfr : float
        SFR surface density in g/cm^2/yr
    """
    Ytot = get_feedback_yield(P, model=Ytot) if isinstance(Ytot, str) else Ytot
    return P / (Ytot * 1.0e5)

def get_scale_height(
    Sigma_gas,
    Sigma_star,
    Omega_d,
    H_star,
    sigma_eff,
    zeta_d=1 / 3.0,
    method="analytic",
    wgas=1,
    wstar=1,
    wdm=1,
):
    """wrapper function to calculate gas scale height either numerically or analytically

    Parameters
    ----------
    Sigma_gas : float or array-like
        gas surface density
    Sigma_star : float or array-like
        stellar surface density
    Omega_d : float or array-like
        galactic rotation speed
    H_star : float or array-like
        stellar scale height
    sigma_eff : str or float or array-like
        velocity dispersion
    zeta_d : float [1/3.]
        parameter for dm and gas distribution
    method : str [analytic or numerical]
        calculate scale height either analytically or numerically

    wgas : int [0 or 1]
        toggle weight term from gas
    wstar : int [0 or 1]
        toggle weight term from stars
    wdm : int [0 or 1]
        toggle weight term from dark matter

    Returns
    -------
    H : float or array-like
        scale height in cm

    Example
    -------
    >>> import prfm
    >>> H=prfm.get_scale_height(Sigma_gas, Sigma_star, Omega_d, H_star, sigma_eff)
    """
    args = (Sigma_gas, Sigma_star, Omega_d, H_star, sigma_eff)
    kwargs = dict(zeta_d=zeta_d, wgas=wgas, wstar=wstar, wdm=wdm)
    w = (wgas << 2) + (wstar << 1) + (wdm << 0)

    if method == "numerical":
        return get_scale_height_numerical(*args, **kwargs)
    else:
        if w == 7:  # 111
            return get_scale_height_analytic(*args, zeta_d=zeta_d)
        elif w == 6:  # 110
            return get_scale_height_star_gas(*args, zeta_d=zeta_d)
        elif w == 5:  # 101
            return get_scale_height_dm_gas(*args, zeta_d=zeta_d)
        elif w == 4:  # 100
            return get_scale_height_gas_only(*args, zeta_d=zeta_d)
        elif w == 3:  # 011
            # no analytic solution provided
            return
        elif w == 2:  # 010
            return get_scale_height_star_only(*args, zeta_d=zeta_d)
        elif w == 1:  # 001
            return get_scale_height_dm_only(*args, zeta_d=zeta_d)
        else:
            # no such gas is expected
            raise ("at least one term must be considered")

def get_scale_height_thick(Sigma_gas, rho_star, sigma_eff, zeta_d=1 / np.pi):
    """approximate solution for H_* >> H_gas without explicitly defining H_*

    Parameters
    ----------
    Sigma_gas : float or array-like
        gas surface density
    rho_star : float or array-like
        stellar(+dark matter) volume density
    sigma_eff : str or float or array-like
        velocity dispersion

    Returns
    -------
    H : float or array-like
        scale height in cm

    Notes
    -----
    Ostriker & Kim (2022) Equation (5)
    """
    g1 = np.pi * _Gconst_cgs * Sigma_gas
    g2 = 32 * np.pi * zeta_d * _Gconst_cgs * rho_star * sigma_eff**2
    return 2 * sigma_eff**2 / (g1 + np.sqrt(g1**2 + g2))

def get_scale_height_thin(Sigma_gas, Sigma_star, sigma_eff):
    """approximate solution for H_* << H_gas without explicitly defining H_*

    Parameters
    ----------
    Sigma_gas : float or array-like
        gas surface density
    Sigma_star : float or array-like
        stellar surface density
    sigma_eff : str or float or array-like
        velocity dispersion

    Returns
    -------
    H : float or array-like
        scale height in cm

    Notes
    -----
    Ostriker & Kim (2022) Equation (5)
    """
    g1 = np.pi * _Gconst_cgs * Sigma_gas
    g2 = 2 * np.pi * _Gconst_cgs * Sigma_star
    return sigma_eff**2 / (g1 + g2)

@np.vectorize
def get_scale_height_analytic(
    Sigma_gas, Sigma_star, Omega_d, H_star, sigma_eff, zeta_d=1 / 3.0
):
    """Analytic solution of the cubic equation for the vertical dynamical equilibirum
    including all three weight terms.

    All inputs must be in c.g.s. units
    """
    args = (Sigma_gas, Sigma_star, Omega_d, H_star, sigma_eff)
    kwargs = dict(zeta_d=zeta_d)
    H_gas_only = get_scale_height_gas_only(*args, **kwargs)
    H_dm_only = get_scale_height_dm_only(*args, **kwargs)

    s_star = Sigma_star / Sigma_gas
    eta_star = H_star / H_gas_only
    eta_dm_sq = (H_dm_only / H_gas_only) ** 2

    a0 = -eta_star * eta_dm_sq
    a1 = (eta_star - 1) * eta_dm_sq
    a2 = (1 + 2 * s_star) * eta_dm_sq + eta_star

    Q = (3 * a1 - a2**2) / 9.0
    R = (9 * a2 * a1 - 27 * a0 - 2 * a2**3) / 54.0

    theta = np.arccos(R / np.sqrt(-(Q**3)))
    h = 2 * np.sqrt(-Q) * np.cos(theta / 3) - a2 / 3.0

    return h * H_gas_only

@np.vectorize
def get_scale_height_numerical(
    Sigma_gas,
    Sigma_star,
    Omega_d,
    H_star,
    sigma_eff,
    zeta_d=1 / 3.0,
    wgas=1,
    wstar=1,
    wdm=1,
    fgas=get_weight_gas,
    fstar=get_weight_star,
    fdm=get_weight_dm,
):
    """Numerical solution of the vertical dynamical equilibrium equation
    including all weight terms.

    All inputs must be in c.g.s. units

    Parameters
    ----------
    fgas : func
    fstar : func
    fdm : func
    """

    def fun(x):
        return (
            get_pressure(Sigma_gas, x, sigma_eff)
            - wgas * fgas(Sigma_gas)
            - wstar * fstar(Sigma_gas, x, Sigma_star, H_star)
            - wdm * fdm(Sigma_gas, x, Omega_d, zeta_d=zeta_d)
        )

    # soln=root(fun,H_gas_init)
    # return soln.x
    return brentq(fun, 1.0e-3 * _pc_cgs, 1.0e6 * _pc_cgs)

def get_self_consistent_solution(
    Sigma_gas,
    Sigma_star,
    Omega_d,
    H_star,
    sigma_eff="tigress_mid",
    zeta_d=1 / 3.0,
    method="analytic",
    tol=1.0e-5,
    niter=16,
):
    """Iteratively solve for sigma_eff(P) to get converged weight, H, and sigma_eff

    Parameters
    ----------
    Sigma_gas: float or array-like
        gas surface density
    Sigma_star: float or array-like
        stellar surface density
    Omega_d: float or array-like
        galactic rotation speed
    H_star: float or array-like
        stellar scale height
    sigma_eff : float or str
        effective vertical velocity dispersion
    zeta_d: float [1/3.]
        parameter for dm and gas distribution
    method: str [analytic or numerical]
        calculate scale height either analytically or numerically

    Returns
    -------
    Wtot : float or array
    H : float or array
    sigma_eff: float or array
    """
    # make initial guess
    if not isinstance(sigma_eff, str):
        sigma0 = sigma_eff
    else:
        if not (sigma_eff in _sigma_eff_models):
            raise IndexError("sigma_eff models: ", list(_sigma_eff_models.keys()))
        sigma0 = 15.0e5
    args = (Sigma_gas, Sigma_star, Omega_d, H_star, sigma0)
    kwargs = dict(zeta_d=zeta_d, method=method)
    Hprev, wgas, wstar, wdm = get_weights(*args, **kwargs)
    wtot_prev = wgas + wstar + wdm

    # return if velocity diseprsion is a constant
    if not isinstance(sigma_eff, str):
        return wtot_prev, Hprev, sigma0 * np.ones_like(Hprev)

    # iterative solve
    for i in range(niter):
        new_sigma_eff = get_sigma_eff(wtot_prev, model=sigma_eff)
        args = (Sigma_gas, Sigma_star, Omega_d, H_star, new_sigma_eff)
        Hnext, wgas, wstar, wdm = get_weights(*args, **kwargs)
        wtot_next = wgas + wstar + wdm
        # L1_norm = np.sum(np.abs(wtot_next/wtot_prev - 1))
        L1_norm = np.sum(np.abs(Hnext / Hprev - 1))
        if L1_norm < tol:
            break
        wtot_prev = np.copy(wtot_next)
        Hprev = np.copy(Hnext)

    return wtot_next, Hnext, get_sigma_eff(wtot_next, model=sigma_eff)

class PRFM(object):
    """a wrapper class for scale height, weight, SFR calculations

    Parameters
    ----------
    Sigma_gas : float, array-like
        gas surface density
    Sigma_star : float, array-like
        stellar surface density
    Omega_d : float, array-like
        galactic rotation speed
    H_star : float, array-like
        stellar scale height
    sigma_eff : str, float, array-like
        velocity dispersion
    zeta_d : float [1/3.]
        parameter for dm and gas distribution
    method : str [analytic or numerical]
        calculate scale height either analytically or numerically
    """

    def __init__(
        self,
        Sigma_gas=None,
        Sigma_star=None,
        H_star=None,
        rho_star=None,
        Omega_d=None,
        rho_dm=None,
        sigma_eff="tigress-mid",
        astro_units=True,
    ):
        # initialize method attribute
        self._method = "analytic"

        # set model name based on sigma_eff
        self.name = sigma_eff if isinstance(sigma_eff, str) else "constant"

        # initialize unit conversion factors
        self.units = self._set_units()

        # setting up gas surface density
        if Sigma_gas is None:
            # gas density must be given
            raise ArgumentError("Sigma_gas must be provided")
        self._Sigma_gas = Sigma_gas

        # setting up stellar disk
        if (Sigma_star is None) and (rho_star is None):
            self._Sigma_star = 0
            self._rho_star = 0
            self._wstar = 0
        else:
            self._wstar = 1
            if H_star is None:
                if rho_star is None:
                    # this is limit for H_star << H_gas
                    self._stellar_disk = "thick"
                    self._Sigma_star = Sigma_star
                elif Sigma_star is None:
                    # this is limit for H_star >> H_gas
                    self._stellar_disk = "thick"
                    self._rho_star = rho_star
                else:
                    # both given
                    self._stellar_disk = "general"
                    self._H_star = Sigma_star / (2 * rho_star)
            else:
                if rho_star is None:
                    # Sigma_star is given
                    rho_star = Sigma_star / (2 * H_star)
                elif Sigma_star is None:
                    # rho_star is given
                    Sigma_star = rho_star * 2 * H_star
                else:
                    # both given
                    raise ArgumentError(
                        "cannot provide all three: Sigma_star, rho_star, H_star"
                    )

                self._stellar_disk = "general"
                self._Sigma_star = Sigma_star
                self._rho_star = rho_star
                self._H_star = H_star
        # unit conversion done before DM part
        if astro_units:
            self._astro_to_cgs()

        # setting up dark matter
        if (Omega_d is None) and (rho_dm is None):
            self._Omega_d = 0
            self._rho_dm = 0
            self._wdm = 0
        else:
            self._wdm = 1
            # this conversion assumes flat rotation
            _fourpiG = 4 * np.pi * ac.G.cgs.value
            if rho_dm is None:
                # Omega_d is given
                if astro_units:
                    Omega_d = Omega_d * self.units["Omega_d"].cgs.value
                rho_dm = Omega_d**2 / _fourpiG
            elif Omega_d is None:
                # rho_dm is given
                if astro_units:
                    rho_dm = rho_dm * self.units["rho_dm"].cgs.value
                Omega_d = np.sqrt(_fourpiG * rho_dm)
            else:
                # both given
                raise ArgumentError("cannot provide both Omega_d and rho_dm")

            self._Omega_d = Omega_d
            self._rho_dm = rho_dm

        if not isinstance(sigma_eff, str):
            if astro_units:
                u = self.units["sigma_eff"].cgs.value
                sigma_eff = sigma_eff * u
            self._sigma_eff = sigma_eff
        else:
            self._sigma_eff_model = sigma_eff

        # set arguments that will be passed to the functions
        if self._stellar_disk == "general":
            self._args = (
                self._Sigma_gas,
                self._Sigma_star,
                self._Omega_d,
                self._H_star,
                sigma_eff,
            )
        elif self._stellar_disk == "thick":
            self._args = (self._Sigma_gas, self._rho_star, sigma_eff)
        elif self._stellar_disk == "thin":
            self._args = (self._Sigma_gas, self._Sigma_star, sigma_eff)
        # store parameters in astro-friendly units
        self._cgs_to_astro()

        # set weight functions
        self._weight_functions = [get_weight_gas, get_weight_star, get_weight_dm]

    def set_weight_function_gas(self, func):
        """Update weight function for gas"""
        self._weight_functions[0] = func
        self._method = "numerical"

    def set_weight_function_star(self, func):
        """Update weight function for star"""
        self._weight_functions[1] = func
        self._method = "numerical"

    def set_weight_function_dm(self, func):
        """Update weight function for star"""
        self._weight_functions[2] = func
        self._method = "numerical"

    def reset_arg_list(self, sigma_eff):
        # reset arguments with new velocity dispersion
        if self._stellar_disk == "general":
            self._args = (
                self._Sigma_gas,
                self._Sigma_star,
                self._Omega_d,
                self._H_star,
                sigma_eff,
            )
        elif self._stellar_disk == "thick":
            self._args = (self._Sigma_gas, self._rho_star, sigma_eff)
        elif self._stellar_disk == "thin":
            self._args = (self._Sigma_gas, self._Sigma_star, sigma_eff)

    def _astro_to_cgs(self):
        for var in [
            "Sigma_gas",
            "Sigma_star",
            "H_star",
            "rho_star",
            "rho_dm",
            "Omega_d",
            "sigma_eff",
        ]:
            if hasattr(self, "_" + var):
                u_cgs = self.units[var].cgs.value
                v = getattr(self, "_" + var) * u_cgs
                setattr(self, "_" + var, v)

    def _cgs_to_astro(self):
        for var in [
            "Sigma_gas",
            "Sigma_star",
            "H_star",
            "rho_star",
            "rho_dm",
            "Omega_d",
            "sigma_eff",
        ]:
            if hasattr(self, "_" + var):
                u_cgs = self.units[var].cgs.value
                v = getattr(self, "_" + var) / u_cgs
                setattr(self, var, v)

    def _set_units(self):
        units = dict()
        units["Sigma_gas"] = units["Sigma_star"] = ac.M_sun / ac.pc**2
        units["rho_star"] = units["rho_dm"] = ac.M_sun / ac.pc**3
        units["Omega_d"] = au.km / au.s / ac.kpc
        units["H_star"] = units["H_gas"] = ac.pc
        units["Wtot"] = ac.k_B / au.cm**3 * au.K
        units["Wgas"] = units["Wstar"] = units["Wdm"] = units["Wtot"]
        units["Ytot"] = 1.0 * au.km / au.s
        units["sigma_eff"] = 1.0 * au.km / au.s
        units["Sigma_SFR"] = ac.M_sun / ac.kpc**2 / au.yr

        return units

    def __repr__(self) -> str:
        args = "PRFM calculator is prepared for\n"
        for var in [
            "Sigma_gas",
            "Sigma_star",
            "rho_star",
            "Omega_d",
            "rho_dm",
            "sigma_eff",
        ]:
            if not hasattr(self, "_" + var):
                continue
            v = np.atleast_1d(getattr(self, var))
            args += "  {}[0]: {}, N={}\n".format(var, v[0], len(v))
        if hasattr(self, "_sigma_eff_model"):
            args += "  sigma_eff_model: {}\n".format(self._sigma_eff_model)
        return args

    def get_scale_height(self, method=None, wgas=1, wstar=1, wdm=1):
        """Wrapper function to get the gas scale height

        Parameters
        ----------
        method : str
            override method if provided
        """

        if method is None:
            method = self._method

        if self._stellar_disk == "general":
            H_gas = get_scale_height(*self._args, wgas=wgas, wstar=wstar, wdm=wdm)
        elif self._stellar_disk == "thick":
            H_gas = get_scale_height_thick(*self._args)
        elif self._stellar_disk == "thin":
            H_gas = get_scale_height_thin(*self._args)

        return H_gas / self.units["H_gas"].cgs.value

    def get_scale_height_numerical(self):
        """Wrapper function to get the gas scale height for generic weight functions"""

        H_gas = get_scale_height_numerical(
            *self._args,
            fgas=self._weight_functions[0],
            fstar=self._weight_functions[1],
            fdm=self._weight_functions[2]
        )
        return H_gas / self.units["H_gas"].cgs.value

    def get_weight_contribution(self):
        """Wrapper function to get weight contribution"""
        if not hasattr(self, "Wtot"):
            self.calc_weights()
        return self.Wgas / self.Wtot, self.Wstar / self.Wtot, self.Wdm / self.Wtot

    def calc_weights(self, method=None):
        """Wrapper function to calculate scale height and all weights

        Parameters
        ----------
        method : str
            override method if provided
        """

        if method is None:
            method = self._method

        if method == "analytic":
            H_gas = self.get_scale_height(wstar=self._wstar, wdm=self._wdm)
        elif method == "numerical":
            H_gas = self.get_scale_height_numerical()

        _H_gas = H_gas * self.units["H_gas"].cgs.value

        self.H_gas = H_gas
        self._H_gas = _H_gas

        W_gas = self._weight_functions[0](self._Sigma_gas)
        if self._stellar_disk == "thick":
            W_star = get_weight_star_thick(*self._args)
        elif self._stellar_disk == "thin":
            W_star = get_weight_star_thin(*self._args)
        else:
            W_star = self._weight_functions[1](
                self._Sigma_gas, _H_gas, self._Sigma_star, self._H_star
            )
        W_dm = self._weight_functions[2](self._Sigma_gas, _H_gas, self._Omega_d)

        for var, v in zip(
            ["H_gas", "Wgas", "Wstar", "Wdm"], [_H_gas, W_gas, W_star, W_dm]
        ):
            u_cgs = self.units[var].cgs.value
            setattr(self, "_" + var, v)  # results in cgs
            setattr(self, var, v / u_cgs)  # results in astro units

        # total weight
        Wtot = W_gas + W_star + W_dm
        u_cgs = self.units["Wtot"].cgs.value
        self._Wtot = Wtot
        self.Wtot = Wtot / u_cgs

        if hasattr(self, "_sigma_eff_model"):
            self._sigma_eff = get_sigma_eff(Wtot, model=self._sigma_eff_model)
            self.sigma_eff = self._sigma_eff / self.units["sigma_eff"].cgs.value

    def calc_self_consistent_solution(self, method=None, niter=16, tol=1.0e-6):
        """Wrapper function to calculate self-consistent solutions

        Update H_gas, W_tot (and all components), sigma_eff
        """
        if method is None:
            method = self._method

        # for model sigma, need to give an initial guess
        if hasattr(self, "_sigma_eff_model"):
            self.reset_arg_list(15.0e5)
        self.calc_weights(method=method)

        # that's it for constant sigma_eff
        # iterative solve for model sigma
        if hasattr(self, "_sigma_eff_model"):
            # iterative solve
            for i in range(niter):
                Hprev = np.copy(self._H_gas)
                self.reset_arg_list(self._sigma_eff)
                self.calc_weights(method=method)
                L1_norm = np.sum(np.abs(self._H_gas / Hprev - 1))
                if L1_norm < tol:
                    break

    def calc_sfr(self, Ytot=1.0e3):
        """Wrapper function to calculate SFR surface density

        Update Sigma_SFR (astro units) and _Sigma_SFR (cgs)
        """
        if not hasattr(self, "_Wtot"):
            self.calc_self_consistent_solution()

        # results in c.g.s.
        Sigma_SFR = get_sfr(self._Wtot, Ytot=Ytot)
        u_cgs = self.units["Sigma_SFR"].cgs.value
        self._Sigma_SFR = Sigma_SFR
        self.Sigma_SFR = Sigma_SFR / u_cgs

    def check_solutions(self):
        self.calc_self_consistent_solution(method="analytic")
        sola = np.copy(self.Wtot)
        self.calc_self_consistent_solution(method="numerical")
        soln = np.copy(self.Wtot)

        L1 = np.mean(np.abs(sola - soln))
        print("difference between analytic and numerical solutions: {}".format(L1))
        return sola, soln


# coding: utf-8
from __future__ import unicode_literals

import shutil
from tempfile import mktemp
import datetime

import mimetypes
from collections import OrderedDict

import tablib
from django.http import HttpResponse
from django.utils.encoding import smart_str, force_bytes

from django.views.generic.list import MultipleObjectMixin

from admin_view.utils.dt import to_aware
from admin_view.views.base import AdminTemplateView

class ReportFormViewMixin(MultipleObjectMixin):
    filter_class = None
    filename = 'Report.xlsx'

    def get_filter_class(self):
        return self.filter_class

    def get_filter(self, queryset=None):
        return self.get_filter_class()(
            self.request.GET or None,
            queryset=queryset or self.get_queryset(),
            view=self
        )

    def is_filtered(self):
        get = self.request.GET.copy()
        get.pop('_export', None)
        return bool(get)

    def get_queryset(self):
        if not self.is_filtered():
            return self.model.objects.none()
        return super().get_queryset()

    def get_columns(self, filter_form):
        fields = filter_form.get_enabled_fields()
        return fields

    def get_export_format(self):
        return self.request.GET.get('_export')

    def get_context_data(self, **kwargs):
        self.object_list = self.model.objects.none()

        context = super().get_context_data(**kwargs)

        filter_form = self.get_filter()
        filter_form.form.is_valid()
        columns = self.get_columns(filter_form)

        export_format = self.get_export_format()

        qs = filter_form.qs
        page = None

        if not export_format:
            paginator, page, qs, is_paginated = self.paginate_queryset(qs, page_size=100)

        export_data = self.get_export_data(qs, columns)

        context['form'] = filter_form
        context['export_data'] = export_data
        context['export_format'] = self.get_export_format()
        context['page'] = page
        return context

    def get_export_data(self, qs, columns):
        export_data = tablib.Dataset(headers=[label for key, label in columns] + ['_instance'])
        for instance in qs:
            export_data.append(self.get_export_data_row(instance, columns) + [instance])
        return export_data

    def get_export_data_row(self, instance, columns):
        row = [self.extract_field(instance, key) for key, label in columns]
        return row

    def extract_field(self, instance, key):
        func = getattr(self, 'extract_' + key, None)
        if func:
            value = func(instance)
        else:
            if hasattr(instance, 'get_%s_display' % key):
                value = getattr(instance, 'get_%s_display' % key)()
            else:
                value = getattr(instance, key)
        if value is None:
            value = ''
        if isinstance(value, datetime.datetime):
            value = to_aware(value).strftime('%Y-%m-%d %H:%M:%S')
        return value

    def get_xlsx_file(self, request, output_file, **kwargs):
        context = self.get_context_data(**kwargs)
        export_format = context['export_format']
        export_dataset = context['export_data']
        del export_dataset[export_dataset.headers[-1]]

        with open(output_file, 'wb') as f:
            f.write(getattr(export_dataset, export_format))

        return

    def get_filename(self):

        return self.filename

    def get(self, request, *args, **kwargs):
        export_format = self.get_export_format()
        if export_format:
            from multiprocessing import Process
            output_file = mktemp(suffix='.' + export_format)
            try:
                p = Process(target=self.get_xlsx_file,
                            args=(request,),
                            kwargs={'output_file': output_file}
                            )
                p.start()
                p.join()
                filename = self.get_filename()
                content_type = mimetypes.guess_type(filename)[0]
                response = HttpResponse(
                    open(output_file, 'rb').read(),
                    content_type=content_type)
                response['Content-Disposition'] = b'attachment; filename="%s"' % force_bytes(
                    filename)
                return response
            finally:
                shutil.rmtree(output_file, ignore_errors=True)

        context = self.get_context_data(**kwargs)
        return self.render_to_response(context)

class ReportFormView(ReportFormViewMixin, AdminTemplateView):
    template_name = "admin/custom_view/report.html"


import tkinter as tk

class Window(tk.Frame):
    def __init__(self, master=None):
        tk.Frame.__init__(self, master)

        self.master = master
        self.init_window()

    def init_window(self):
        self.master.title("GUI")
        self.pack(fill=tk.BOTH, expand=1)
        quitButton = tk.Button(self, text="Quit")
        quitButton.place(x=0, y=0)

root = tk.Tk()
root.geometry("400x300")
app = Window(root)
root.mainloop()


team = [
    'john', 10000,
    'neo', 100,
    'tak', 40500
]

team[2:4] #neo와 100이 튀어나온다

type(team[2:4]) #타입이 list가 튀어나온다
team[2:3] # ['neo'] 가 나옴
team[2] # 'neo' 만 나옴

new_member = ['js', 10]

team = team + new_member #list끼리 덧셈연산 가능하다, 우측 결과값을 왼쪽에 넣는다
#team 은 이렇게 나옴 : ['john', 10000, 'neo', 100, 'tak', 40500, 'js', 10]

#list끼리 덧셈연산하면 list가 그 안으로 뒤에서 합쳐짐
team += new_member #위 team = team + new_member와 같은 코드다

del(team[2]) # 인덱스 2위치의 'neo'가 날아감 삭제!
del(team[2:4])#인덱스 2부터 3까지 날아감


import cv2
import numpy as np
import uuid

im = cv2.imread("test.jpg")
print(im.shape)
print(im[0,0])
print(im[0,0][0])
print(im[0,0][1])
print(im[0,0][2])

x1 = None
y1 = None
x2 = None
y2 = None
backgroundBgr = [255, 255, 255]

for y in range(0, im.shape[0] - 1):
	#print(y)
	for x in range(0, im.shape[1] - 1):
		#print(x)
		bgr = im[y, x]
		for i,v in enumerate(bgr):
			if abs(v - backgroundBgr[i]) > 10:
				if (x1 is None) or (x < x1):
					x1 = x
				if (x2 is None) or (x > x2):
					x2 = x
				if (y1 is None) or (y < y1):
					y1 = y
				if (y2 is None) or (y > y2):
					y2 = y
				break;

print(x1, y1, x2, y2)

cv2.rectangle(im, (x1, y1), (x2, y2), (0,0,255), 3);
cv2.imwrite(str(uuid.uuid4()) + ".jpg", im)

# bgr
'''
cv2.rectangle(im, (10, 10), (20, 20), (0,0,255));

im2 = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)

cv2.imshow("Over the Clouds", im)
cv2.imshow("Over the Clouds - gray", im2)

cv2.waitKey(0)
cv2.destroyAllWindows()
'''


# -*- encoding: utf-8 -*-
from django.conf.urls import patterns, url
from rest_framework.routers import DefaultRouter
from .resources import (
    ServiceViewSet, AdvertViewSet, BookingViewSet
)

router = DefaultRouter()

router.register(r'service', ServiceViewSet)
router.register(r'advert', AdvertViewSet)
router.register(r'booking', BookingViewSet)

urlpatterns = patterns('',
)
urlpatterns += router.urls


from collections.abc import Mapping

from habutax.inputs import StringInput, BooleanInput, IntegerInput, FloatInput, EnumInput, SSNInput
from habutax.fields import StringField, BooleanField, IntegerField, FloatField, EnumField

class Form(object):
    def __init__(self,
                 child_cls,
                 inputs,
                 required_fields,
                 optional_fields,
                 instance=None,
                 solver=None):

        self._name = child_cls.form_name
        assert("." not in self._name)
        self._tax_year = child_cls.tax_year
        self._inputs = inputs
        self._required_fields = required_fields
        self._optional_fields = optional_fields
        self._instance = instance
        self._solver = solver

        for i in self._inputs:
            i.__form_init__(self)
        for f in self._required_fields + self._optional_fields:
            f.__form_init__(self)

    def name(self):
        if self._instance is None:
            return self._name
        else:
            return f'{self._name}:{self._instance}'

    def solver(self):
        assert(self._solver)
        return self._solver

    def inputs(self):
        return self._inputs

    def required_fields(self):
        return self._required_fields

    def fields(self):
        return self.required_fields() + self._optional_fields

class InputForm(Form):
    """Convenience class to create a form which creates fields for each
    input"""
    def __init__(self,
                 child_cls,
                 inputs,
                 **kwargs):

        fields = []
        for i in inputs:
            base_name = i.base_name()
            fn = lambda s, i, v, base_name=base_name: i[base_name]

            if type(i) in [StringInput, SSNInput]:
                fields.append(StringField(base_name, fn))
            elif type(i) is BooleanInput:
                fields.append(BooleanField(base_name, fn))
            elif type(i) is IntegerInput:
                fields.append(IntegerField(base_name, fn))
            elif type(i) is FloatInput:
                fields.append(FloatField(base_name, fn))
            elif type(i) is EnumInput:
                fields.append(EnumField(base_name, i.enum, fn))
            else:
                raise TypeError(f'Unexpected input type in InputForm: {type(i)}')

        super().__init__(child_cls, inputs, fields, [], **kwargs)

class FormAccessor(Mapping):
    def __init__(self, mapping, form):
        self.form = form
        self.mapping = mapping

    def __getitem__(self, key):
        if "." not in key:
            key = f'{self.form.name()}.{key}'
        return self.mapping[key]

    def __iter__(self):
        return iter(self.mapping)

    def __len__(self):
        return len(self.mapping)


"""
    Author: Nikhil Patil
    Project: ClassAttender
    Purpose:
        We love efficiency. Being a software engineer and someone who likes to have fun, I created this project
    to allow me to sleep in or play games until the last minute possible. And also to ensure that I don't forget to
    attend my class sessions.

    Code:
        1) Read the classes.txt file to find the time and day of my classes. Along with the link and any other details.
        2) Create a batch file that will be executed by the Task Scheduler to run this script for Windows.
        3) Create tasks in the Task Scheduler that runs the batch script whenever one of your sessions is about to begin
        4) When executed, this script will check if any of your classes are about to begin in 5 minutes or began 5
           minutes ago. If so, the script will launch the site and join the live session for you.
"""

import os
import subprocess
from datetime import datetime
import cv2
import pyautogui
import webbrowser
from time import sleep
import platform

# Define the path to chrome exe
browser_path = r'C:/Program Files (x86)/Google/Chrome/Application/chrome.exe %s'
# Path to python
python_path = r'C:\ProgramData\Anaconda3\python.exe'

# Launches the session
def openClass(info):
    # For elearning collaborate session
    if 'elearning' in info[0]:
        # Open the elearning page
        webbrowser.get(browser_path).open('https://elearning.utdallas.edu')

        sleep(3)  # Wait for page to load. Sleep for 3 seconds
        # Click Login if on the login screen
        findAndClick('elearningLogin.png')

        sleep(2)  # Wait for authentication
        # Launch the session link
        webbrowser.get(browser_path).open(info[0])

        # For the professors that created Regular class sessions
        if 'Regular' in info:
            sleep(5)
            # Find and click on the regular class icon
            findAndClick('regularClass.png')
        else:
            sleep(5)
            # Find the course room button
            findAndClick('courseRoom.png')
            sleep(1)
            # Find the join room button
            findAndClick('joinCourseRoom.png')
    # For webex meetings
    elif 'webex' in info[0]:
        # Go to the link
        webbrowser.get(browser_path).open(info[0])
        sleep(10)  # Wait for WebEx web client to load
        # Find the join meeting button and start session
        findAndClick('webexJoinMeeting.png')

# Takes a screenshot of the screen, find the specified sub-image in the screenshot, click it
def findAndClick(findSS=""):
    screen = pyautogui.screenshot()
    # Save the screenshot as screenshot.png
    screen.save(r'screenshot.png')
    # Creates a grayscale image
    img = cv2.imread('screenshot.png', 0)
    template = cv2.imread(findSS, 0)

    # Finds the width and height of the template/button
    w, h = template.shape[::-1]

    # Specify the template being used
    method = cv2.TM_SQDIFF_NORMED

    # Apply template Matching
    res = cv2.matchTemplate(img, template, method)

    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)

    # Used as a threshold measure. Smaller the number the better match it is
    if min_val > 0.01:
        return

    # Find the top left of the image
    top_left = min_loc
    # Find the center of the button
    button_center = (int(top_left[0] + w / 2), int(top_left[1] + h / 2))

    # print(button_center) # Prints the (x, y) coordinate of the center

    # Cute little animation of moving the mouse to the button
    pyautogui.moveTo(*button_center, duration=0.5)
    # Press the button
    pyautogui.click(*button_center)

# Get the class schedule
def getClasses(filePath="classes.txt"):
    data = []
    with open(filePath, "r") as file:
        for line in file.readlines():
            line.replace('\n', '')
            info = line.split(", ")

            # Allow some comments in the classes.txt file
            if '#' in info[0] or '//' in info[0] or len(info) < 4:
                continue

            # Get the class name. Used while creating task later
            className = info[0]
            # Store the time on 24 hr clock
            classTime = datetime.strptime(info[1], '%I:%M %p')

            # Create a list of 3-letter day names, in uppercase
            classDays = info[2].split()
            temp = []
            for day in classDays:
                temp.append(day.upper())
            classDays = temp

            # Append info to data
            data.append([className, classTime, classDays, *info[3:]])
    return data

# Create tasks in the Task Scheduler
def createTasks(data, batPath="{}\\classAttender.bat".format(os.getcwd())):
    if batPath is None or len(data) == 0:
        return

    for item in data:
        # Get the days from the list without the [] or ''
        days = ""
        for day in item[2]:
            days += "{},".format(day)
        days = days[:-1]

        # Query to create a Task in Task Scheduler with the name of the class for each class
        # Repeats every week on specified date and time
        cmd = r"schtasks /create /tn {}ClassAttender /tr {} /sc weekly /d {} /st {:02d}:{:02d} /f"\
            .format(item[0], batPath, days, item[1].hour, item[1].minute)

        # Execute the command to create the task
        os.system(cmd)

# Create a batch file that will be executed by the task scheduler
def createBatch(batName="classAttender.bat"):
    # Changes directory to current directory
    # Executes the python script using python path provided above
    # Exit the batch script
    with open(batName, 'w') as file:
        file.writelines([r'cd {}'.format(os.getcwd()), '\n',
                         r'{} {}'.format(python_path, __file__), '\n',
                         r'exit'])
    # Return the absolute path to the batch file
    return "{}\\{}".format(os.getcwd(), batName)

# Deletes any previously scheduled tasks
def deleteTasks():
    # Find all tasks with ClassAttender in their names
    for line in subprocess.Popen("schtasks | findstr ClassAttender",
                                 shell=True, stdout=subprocess.PIPE).stdout.readlines():
        line = line.decode("utf-8")
        # Extract the class name
        className = line.split('ClassAttender')[0]
        if len(className) > 0:
            # Remove the task
            cmd = r"schtasks /delete /tn {}ClassAttender /f".format(className)
            os.system(cmd)

def main():
    # Support for other OS
    otherOS = False
    # Stores the minute difference to the next class. 1440 minutes in a day
    nextClass = 1441

    # Get the current time
    currTime = datetime.now()
    # Find today: MON, TUE, WED, THU, FRI, SAT, SUN
    today = currTime.strftime("%a").upper()

    # Reads the classes.txt file and creates a dictionary with the data
    data = getClasses()

    # Check if the current OS is Windows
    if 'Windows' in platform.system():
        # Check if  the batch file exists or not
        if os.path.exists("classAttender.bat") is False:
            # Create the batch file
            batPath = createBatch()
            # Delete previous tasks
            deleteTasks()
            # Create tasks that start this script whenever a session is about to start
            createTasks(data, batPath)
    else:
        # If the script is running on non-windows machine
        otherOS = True

    for item in data:
        classTime = item[1]

        # Find the difference between current time and when the session begins
        hourDiff = (classTime.time().hour - currTime.time().hour)
        minDiff = (classTime.time().minute - currTime.time().minute)
        timeDiff = abs(hourDiff * 60 + minDiff)
        # print(classTime.time(), currTime.time(), timeDiff)

        # If the session begins in or began 5 minutes ago, launch the session
        if timeDiff <= 5 and today in item[2]:
            openClass(item[3:])
        # If on non-windows OS and did not launch the class
        elif otherOS:
            # Add the difference to the next closest class
            if timeDiff < nextClass:
                nextClass = timeDiff

    # Instead of having tasks for other OS, the script will continuously run
    if otherOS:
        # Sleep until next class is about to start
        sleep(nextClass)
        # Run the script again
        main()

if __name__ == "__main__":
    main()


# Set some variables to make this easier next time
import os

SOURCEDIR = os.environ['GUIDIR']
QTDIR = 'C:/qt/4.5.2/bin'
PYINSTALLER = 'C:/apps/pyinstaller'

# Process the includes and excludes first

data_files = [
            ('README.TXT', '%s/README.TXT' % SOURCEDIR, 'DATA'),
            ('license.txt', '%s/license.txt' % SOURCEDIR, 'DATA'),
            ('openvpn-client.ico', '%s/ui/resources/openvpn-client.ico' % SOURCEDIR, 'DATA')
            ]

includes = []
excludes = ['_gtkagg', '_tkagg', 'bsddb', 'curses', 'email', 'pywin.debugger',
            'pywin.debugger.dbgcon', 'pywin.dialogs', 'tcl',
            'Tkconstants', 'Tkinter']
packages = []
dll_excludes = []
dll_includes = [('QtCore4.dll', '%s/QtCore4.dll' % QTDIR,
                'BINARY'), ('QtGui4.dll', '%s/QtGui4.dll' % QTDIR,
                'BINARY'), ('QtNetwork4.dll', '%s/QtNetwork4.dll' % QTDIR,
                'BINARY')]

# Set up the more obscure PyInstaller runtime options

options = [('O', '', 'OPTION')]

guianalysis = Analysis(['%s/support/_mountzlib.py' % PYINSTALLER,
           '%s/support/useUnicode.py' % PYINSTALLER,
           '%s/opensesame.py' % SOURCEDIR,],
                    pathex=[],
                    hookspath=[],
                    excludes=excludes)

guipyz = PYZ(guianalysis.pure, level=9)

tapanalysis = Analysis(['%s/support/_mountzlib.py' % PYINSTALLER,
           '%s/support/useUnicode.py' % PYINSTALLER,
           '%s/installer/tapinstaller.py' % SOURCEDIR],
                    pathex=[],
                    hookspath=[],
                    excludes=excludes)

tappyz = PYZ(tapanalysis.pure, level=9)

guiexecutable = EXE(guipyz,
                 guianalysis.scripts + includes + packages + options,
                 exclude_binaries=1,
                 name="OpenSesame.exe",
                 debug=False,
                 console=False,
                 strip=False,
                 upx=False,
                 icon='%s/ui/resources/openvpn-client.ico' % SOURCEDIR,
                 version='%s/installer/gui-versioninfo.txt' % SOURCEDIR )

tapexecutable = EXE(tappyz,
                 tapanalysis.scripts + includes + packages + options,
                 exclude_binaries=1,
                 name="tapinstaller.exe",
                 debug=False,
                 console=True,
                 strip=False,
                 upx=False,
                 icon=None,
                 version='%s/installer/tapinstaller-versioninfo.txt' % SOURCEDIR )

collect = COLLECT( guiexecutable, tapexecutable,
                  guianalysis.binaries + tapanalysis.binaries - dll_excludes + dll_includes + data_files,
                  name="dist",
                  strip=False,
                  upx=False)


"""Provides tests for the target creator module."""

# Authors: Federico Raimondo <f.raimondo@fz-juelich.de>
#          Sami Hamdan <s.hamdan@fz-juelich.de>
# License: AGPL

from sklearn.preprocessing import MinMaxScaler, StandardScaler

from julearn.pipeline import JuTargetPipeline, TargetPipelineCreator
from julearn.transformers.target import TargetConfoundRemover

def test_TargetPipelineCreator() -> None:
    """Test the target pipeline creator."""

    creator = TargetPipelineCreator()
    creator.add("zscore")
    creator.add("scaler_minmax")
    creator.add("confound_removal", confounds="confounds")
    pipeline = creator.to_pipeline()

    assert isinstance(pipeline, JuTargetPipeline)
    assert len(pipeline.steps) == 3
    assert isinstance(pipeline.steps[0][1], StandardScaler)
    assert isinstance(pipeline.steps[1][1], MinMaxScaler)
    assert isinstance(pipeline.steps[2][1], TargetConfoundRemover)

def test_TargetPipelineCreator_repeated_names() -> None:
    """Test the target pipeline creator."""

    creator = TargetPipelineCreator()
    creator.add("zscore")
    creator.add("zscore")
    pipeline = creator.to_pipeline()

    assert isinstance(pipeline, JuTargetPipeline)
    assert len(pipeline.steps) == 2
    assert pipeline.steps[0][0] == "zscore"
    assert pipeline.steps[1][0] == "zscore_1"