import logging
import os
from copy import deepcopy

from harmony.util import shortened_id
from harmony.repository_state import RepositoryState

logger = logging.getLogger(__name__)

def commit(local_location_id, working_directory, location_states, repository_state):
    """
    Scan the given working directory for changes and commit them to local
    state storage.
    That is, update location_states[local_location_id] with the
    new current file states (digests, "who has what?").
    Also update repository_state info ("who made which content decision in what
    order?")

    Parameters:

    local_location_id:
        ID of the location that is considered local (i.e. the one that belongs
        to the working_directory instance)

    working_directory:
        WorkingDirectory instance representing the local working directory.

    location_states:
        LocationStates instance representing the local location state storage.
        Will (possibly) be modified.

    repository_state:
        RepositoryState instance representing the local repository state
        storage. Will (possibly) be modified.

    return:
        True iff any change was recorded.
    """

    id_ = local_location_id
    short_id = shortened_id(id_)

    paths = set(working_directory.get_filenames()) \
            | set(location_states.get_all_paths(id_))

    # 1. update location state
    #    - detect renames (add WIPE entries later for those)
    #    - when a file is *added* that is known to other locations w/
    #      different digest, let user confirm what he wants to do (see
    #      above)
    #    - increase local clock
    #
    # 2. update repository state
    #    - if file changed in step 1:
    #      clock = current clock for local + max for each other location
    #      hash = current local hash
    #      (deviate from this if user selected to do something else)
    #    - if file did not change:
    #      no change in hash or clock

    # Do all the file scanning before so we can be sure to do it at most
    # once per file in the WD
    wd_states = {
        path: working_directory.generate_file_state(path)
        for path in paths
        if working_directory.file_maybe_modified(
            location_states.get_file_state(id_, path)
        )
    }

    location_state_cache = {
        path: location_states.get_file_state(id_, path)
        for path in paths
    }

    any_change = False
    for path in paths:

        if path in wd_states:
            file_state = location_state_cache[path]
            new_file_state = wd_states[path]
            changed = location_states.update_file_state(id_, new_file_state)
            if changed:
                any_change = True

                # If the file vanished but a new one with the same digest
                # popped up, consider that a rename.
                # Rename means, the old file is WIPEd (instead of just
                # locally removed) and the new file is added as usual
                if not new_file_state.exists():
                    logger.debug('{} vanished'.format(new_file_state.path))

                    # Iterate over paths to find a possible rename target
                    for path2 in paths:
                        # Rename to itself does not make sense
                        # Rename to a file that has not changed (or better: just appeared) does not make sense
                        if path2 == path or path2 not in wd_states:
                            continue

                        path2_state = location_state_cache[path2]
                        new_path2_state = wd_states[path2]
                        logger.debug('{} rename candidate {} ex before={} ex now={} self.digest={} candidate.digest={}'.format(
                            path, path2, path2_state.exists(),
                            new_path2_state.exists(),
                            file_state.digest, new_path2_state.digest
                        ))

                        if not path2_state.exists() \
                           and new_path2_state.exists() \
                           and new_path2_state.digest == file_state.digest:
                            logger.info('Detected rename: {} -> {}'.format(path, path2))
                            new_file_state.wipe = True
                            new_file_state.digest = file_state.digest
                            break

                repository_state.update_file_state(
                    new_file_state,
                    id_,
                    location_states.get_clock(id_) + 1,
                )
                logger.debug('{} committed: {} clk={}'.format(short_id, new_file_state.path, location_states.get_clock(id_) + 1))
            else:
                logger.debug('{} not actually changed: {}'.format(short_id, path))
        else:
            logger.debug('{} not changed: {}'.format(short_id, path))

    return any_change

def merge(local_state, remote_state, merger_id):
    """
    Merge two repository states ('local' and 'remote') into a common state if
    possible, auto-detecting if a change only happened on one side and
    propagating those changes.
    For cases in which a file was changed on both sides, return details of the
    conflict.

    local_state:
        RepositoryState() instance that reflects the local repository state.
    remote_state:
        RepositoryState() instance that reflects the remote repository state.
    merger_id:
        ID of the repository conducting the merge (assumed to correspond
        to the 'local' repository)
    return:
        A pair (conflicts, merged).
        $conflicts is a dictonary of the form { path: (local_entry, remote_entry),
        ... } whereas $path denotes the path of a file in conflict and $local_entry
        and $remote_entry refer to the RepositoryState.Entry instances for that
        file that are in conflict.
        $merged is a newly created RepositoryState instance with selected merged
        repository states.
        If $conflicts is empty, $merged covers all files present either locally or
        remotely.
    """
    local_paths = set(local_state.get_paths())
    remote_paths = set(remote_state.get_paths())

    merged = RepositoryState(None)
    conflicts = {}

    for p in local_paths - remote_paths:
        merged[p] = local_state[p]

    for p in remote_paths - local_paths:
        merged[p] = remote_state[p]

    # conflicts can only arise in paths that are specified in both state
    # files
    paths = set(local_state.get_paths()) & set(remote_state.get_paths())

    for path in paths:
        local = local_state[path]
        remote = remote_state[path]

        c = local.clock.compare(remote.clock)
        if c is None:
            if local.contents_different(remote):
                logger.debug('merge: {} in conflict: {} <-> {}'.format(
                    path, local.clock, remote.clock
                ))
                conflicts[path] = (local, remote)
            else:
                logger.debug('merge: {} automerged (same content)'.format(path))
                m = deepcopy(local)
                m.clock.update(remote.clock)
                m.clock.increase(merger_id)
                merged[path] = m

        elif c < 0:
            logger.debug('merge: {} newer on remote'.format(path))
            merged[path] = remote

        else: # c >= 0:
            logger.debug('merge: {} same version or newer on local'.format(path))
            merged[path] = local

    return conflicts, merged

def auto_rename(working_directory, repository_state):
    from harmony.working_directory import WorkingDirectory

    assert isinstance(working_directory, WorkingDirectory)
    """
    Apply automatic renaming in the given working_directory.
    That is, if working dir contains files that are WIPEd in $repository_state but
    are present under a different name, automatically rename those to obtain
    the repository file at a low cost.

    Repository.commit() should be called after calling this to commit the
    changes to the working directory.

    precondition: WD clean
    """

    # Automatically apply auto-renaming
    # Auto-renaming
    # -------------
    # 1. Find any files $A with a WIPE entry.
    # 2. Compute/get their digest (from location state)
    # 3. Find a non-wiped file $B in repo that does not exist in the WD
    # 4. Rename $A to $B

    for path, entry in repository_state.files.items():
        logger.debug('auto_rename: {}: path={} wipe={} in_wd={}'.format(path, entry.path, entry.wipe, (entry.path in working_directory)))
        if entry.wipe and (entry.path in working_directory):
            possible_targets = {
                e.path for e in repository_state.files.values()
                if e.path != path and e.digest == entry.digest and not e.wipe
            }
            logger.info(
                '{} could be auto-renamed to any of {}'.format(
                    path, possible_targets
                )
            )
            if possible_targets:
                (working_directory.path / path).rename(working_directory.path / possible_targets.pop())


import numpy as np

import jax
from jax import lax, random, numpy as jnp

import flax
from flax.core import freeze, unfreeze
from flax import linen as nn
from flax import optim

from typing import Any, Callable, Sequence, Optional
import pickle

from tensorflow import keras

file_prefix = "struct"

activation = nn.relu # activation function

M = 300  # width parameter
L = 20  # depth

alpha = 10e-5 # learning rate
epochs = 3000
kernel_steps = [0,1,10,20,50,100,200,300,400,500,600,700,800,900,1000,1200,1500,1700,2000,2500,3000] # epochs at which the NTK is computed

var_w_s = [1.0,2.0,2.2]  # variance parameter \sigma_w^2
var_b = 0.               # variance parameter \sigma_b^2

# custom fully-connected network (MLP) class
class MLP(nn.Module):
    widths: Sequence[int] # We need to specify all the layer width (including input and output widths)
    v_w: float # variance parameter \sigma_w^2
    v_b: float # variance parameter \sigma_b^2

    activation: Callable # activation function (the same in all the hidden layers)
    kernel_init: Callable = jax.nn.initializers.normal # Gaussian initialization
    bias_init: Callable = jax.nn.initializers.normal

    def setup(self):
        self.layers = [nn.Dense(self.widths[l+1],
                               kernel_init = self.kernel_init(jnp.sqrt(self.v_w/self.widths[l])),
                               bias_init = self.bias_init(jnp.sqrt(self.v_b))
                               ) for l in range(len(self.widths)-1)]

    def __call__(self, inputs):
        x = inputs
        for i, lyr in enumerate(self.layers[:-1]):
            x = lyr(x)
            x = self.activation(x)
        x = self.layers[-1](x)
        return x

# the NTK on a single pair of samples (x1,x2)
def K(model):

    def K(x1,x2,params):
        f1 = jax.jacobian(lambda p: model.apply(p,x1))(params)
        f2 = jax.jacobian(lambda p: model.apply(p,x2))(params)
        leaves, struct = jax.tree_util.tree_flatten(jax.tree_multimap(jnp.multiply,f1,f2))

        return sum([jnp.sum(leaf) for leaf in leaves])

    return jax.jit(K)

# the NTK matrix (vectorization of K)
def K_matr(model):
    _K = K(model)

    def K_matr(X,Y,params):
        f = lambda x1,x2: _K(x1,x2,params)
        return jax.vmap(jax.vmap(f,(None,0)),(0,None))(X,Y)

    return jax.jit(K_matr)

# MSE loss function
def mse(x_batched, y_batched):
    def mse(params):
        # MSE on a single pair (x,y)
        def squared_error(x, y):
            pred = model.apply(params, x)
            return jnp.inner(y-pred, y-pred)/2.0
        return jnp.mean(jax.vmap(squared_error)(x_batched,y_batched), axis=0) #Vectorized MSE
    return jax.jit(mse)

 # Load and preprocess MNIST

n_class = 10
ker_size_per_class = 10
mnist_n0 = 28*28

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.reshape(x_train.shape[0], mnist_n0)
x_test = x_test.reshape(x_test.shape[0], mnist_n0)

# choose subset of data with ker_size_per_class samples from each class
ind = []
for k in range(ker_size_per_class):
    ind += list(np.random.choice(np.argwhere(y_train==k).ravel(), size=ker_size_per_class, replace=False))

x_train, x_test = x_train/255.,x_test/255.
y_train, y_test = keras.utils.to_categorical(y_train, 10), keras.utils.to_categorical(y_test, 10)

x_ker = x_train[ind] # We compute the NTK only on a subset of samples
y_ker = y_train[ind]

# -------

key = random.PRNGKey(0)
subkeys = jax.random.split(key, num=len(var_w_s))

widths = [mnist_n0]+[M]*L+[n_class]
optimizer_def = optim.Adam(learning_rate=alpha) # Define Adam optimizer

loss = mse(x_train, y_train) # train loss function
loss_grad_fn = jax.value_and_grad(loss) # function to get loss value and gradient
test_loss = mse(x_test, y_test) # test loss function

for var_w, subkey in zip(var_w_s, subkeys):

    model = MLP(widths = widths, v_w=var_w, v_b=var_b, activation = activation) # Define MLP model
    params = model.init(subkey, x_train) # Initialize model
    optimizer = optimizer_def.create(params) # Create optimizer with initial parameters

    K_t = []
    loss_t = []
    test_loss_t = []

    K_func = K_matr(model)

    for i in range(epochs+1):
        loss_val, grad = loss_grad_fn(optimizer.target) # Get gradient and train loss value
        test_loss_val = test_loss(optimizer.target)  # Get test loss value

        test_loss_t.append(test_loss_val)
        loss_t.append(loss_val)

        # Compute the NTK for the chosen epochs
        if i in kernel_steps:
            print('Loss step {}: '.format(i), loss_val, test_loss_val)
            K_t.append(K_func(x_ker,x_ker,optimizer.target))

        optimizer = optimizer.apply_gradient(grad) # Update optimizer parameters

    # Save the results
    pickle.dump(jnp.array(K_t), open( "ntk_dynamics/"+file_prefix+"_w"+str(int(var_w*10))+"M"+str(M)+"L"+str(L), "wb" ) )
    pickle.dump(jnp.array(loss_t), open( "ntk_dynamics/"+file_prefix+"_loss_w"+str(int(var_w*10))+"M"+str(M)+"L"+str(L), "wb" ) )
    pickle.dump(jnp.array(test_loss_t), open( "ntk_dynamics/"+file_prefix+"_test_loss_w"+str(int(var_w*10))+"M"+str(M)+"L"+str(L), "wb" ) )


import wttest
from wiredtiger import stat
from wtscenario import make_scenarios

# Test compact behaviour with overflow values.
class test_compact03(wttest.WiredTigerTestCase):

    uri='table:test_compact03'

    fileConfig = [
        ('1KB', dict(fileConfig='allocation_size=1KB,leaf_page_max=1KB')),
        ('4KB', dict(fileConfig='allocation_size=4KB,leaf_page_max=4KB')),
    ]

    useTruncate = [
        ('no_truncate', dict(truncate=False)),
        ('truncate', dict(truncate=True))
    ]
    scenarios = make_scenarios(fileConfig, useTruncate)

    # Enable stats and use a cache size that can fit table in the memory.
    conn_config = 'statistics=(all),cache_size=100MB'

    # We want to test how compaction interacts with overflow values. Test flow is as follows:
    #
    # 1. Populate a table with relatively small page size.
    # 2. Checkpoint and get stats on the table to confirm the size.
    # 3. Add few thousand overflow values. It is expected that these will be written at the end of
    #    file.
    # 4. Perform checkpoint to ensure overflow values are written on disk.
    # 5. Delete middle ~90% of the normal values in the table.
    # 6. Perform checkpoint so compact can find something to work with.
    # 7. Call compact.
    # 8. Get stats on compacted table expecting that there will no change in size given we have
    #    overflow keys at the end of file.
    # 9. Insert some normal values again. They will be written in the free extents in the middle
    #    of the file. Therefore, expect no increase in file size.
    #
    # We want to have around 20000 leaf pages. With minimum 1KB page allocation size, the table
    # is expected to have at least 25 MByte worth of data. We can then experiment with deleting
    # range of keys in middle to test how compaction works.

    normalValue = "abcde" * 10
    overflowValue = "abcde" * 1000
    nrecords = 400000 # To create ~25 MB table
    expectedTableSize = 20 # Mbytes
    nOverflowRecords = 5000

    # Return stats that track the progress of compaction.
    def getCompactProgressStats(self):
        cstat = self.session.open_cursor(
            'statistics:' + self.uri, None, 'statistics=(all)')
        statDict = {}
        statDict["pages_reviewed"] = cstat[stat.dsrc.btree_compact_pages_reviewed][2]
        statDict["pages_skipped"] = cstat[stat.dsrc.btree_compact_pages_skipped][2]
        statDict["pages_rewritten"] = cstat[stat.dsrc.btree_compact_pages_rewritten][2]
        cstat.close()
        return statDict

    # Return the size of the file
    def getSize(self):
        # To allow this to work on systems without ftruncate,
        # get the portion of the file allocated, via 'statistics=(all)',
        # not the physical file size, via 'statistics=(size)'.
        cstat = self.session.open_cursor(
            'statistics:' + self.uri, None, 'statistics=(all)')
        sz = cstat[stat.dsrc.block_size][2]
        cstat.close()
        return sz

    # Create a table, add keys with both big and small values.
    def test_compact03(self):
        # FIXME-WT-11399: check the other assertions that are skipped when the tiered hook is
        # enabled.
        if self.runningHook('tiered'):
            self.skipTest("this test generates occasional rollback errors when tiered is enabled")

        mb = 1024 * 1024
        # 1. Create a table with relatively small page size.
        params = 'key_format=i,value_format=S,' + self.fileConfig
        self.session.create(self.uri, params)
        c = self.session.open_cursor(self.uri, None)
        for i in range(self.nrecords):
            c[i] = self.normalValue
        c.close()

        # 2. Checkpoint and get stats on the table to confirm the size.
        self.session.checkpoint()
        sizeWithoutOverflow = self.getSize()
        self.pr('After populate ' + str(sizeWithoutOverflow // mb) + 'MB')
        if not self.runningHook('tiered'):
            self.assertGreater(sizeWithoutOverflow, self.expectedTableSize * mb)

        # 3. Add overflow values.
        c = self.session.open_cursor(self.uri, None)
        for i in range(self.nOverflowRecords):
            c[i + self.nrecords] = self.overflowValue
        c.close()

        # 4. Perform checkpoint to ensure overflow values are written to the disk.
        self.session.checkpoint()
        sizeWithOverflow = self.getSize()
        self.pr('After inserting overflow values ' + str(sizeWithoutOverflow // mb) + 'MB')
        if not self.runningHook('tiered'):
            self.assertGreater(sizeWithOverflow, sizeWithoutOverflow)

        # 5. Delete middle ~90% of the normal values in the table.
        if self.truncate:
            c1 = self.session.open_cursor(self.uri, None)
            c2 = self.session.open_cursor(self.uri, None)
            c1.set_key((self.nrecords // 100) * 5)
            c2.set_key((self.nrecords // 100) * 95)
            self.assertEqual(self.session.truncate(None, c1, c2, None), 0)
            c1.close()
            c2.close()
        else:
            c = self.session.open_cursor(self.uri, None)
            for i in range((self.nrecords // 100) * 5, (self.nrecords // 100) * 95):
                c.set_key(i)
                self.assertEqual(c.remove(), 0)
            c.close()

        # 6. Perform checkpoint to ensure we have blocks available in the middle of the file.
        self.session.checkpoint()

        # 7 & 8. Call compact. We expect that the overflow values at the end of the file are not
        #        rewritten and therefore the file size will mostly remain the same. Give a leeway
        #        of 10%.
        self.session.compact(self.uri)
        sizeAfterCompact = self.getSize()
        self.pr('After deleting values and compactions ' + str(sizeAfterCompact // mb) + 'MB')
        if not self.runningHook('tiered'):
            self.assertGreater(sizeAfterCompact, (sizeWithOverflow // 10) * 9)

            # Verify that we did indeed rewrote some pages but that didn't help with the file size.
            statDict = self.getCompactProgressStats()
            self.assertGreater(statDict["pages_reviewed"],0)
            self.assertGreater(statDict["pages_rewritten"],0)
            self.assertEqual(statDict["pages_rewritten"] + statDict["pages_skipped"],
                            statDict["pages_reviewed"])

        # 9. Insert some normal values and expect that file size won't increase as free extents
        #    in the middle of the file will be used to write new data.

        # Insert around ~50% of the normal values in the table that we deleted earlier.
        c = self.session.open_cursor(self.uri, None)
        for i in range((self.nrecords // 100) * 25, (self.nrecords // 100) * 75):
            c.set_key(i)
            c.set_value(self.normalValue)
            self.assertEqual(c.update(),0)
        c.close()

        # Perform compact.
        self.session.compact(self.uri)

        # Test that the file size doesn't increase.
        sizeAfterNewInserts = self.getSize()
        self.pr('After Inserting bunch of values ' + str(sizeAfterNewInserts // mb) + 'MB')
        if not self.runningHook('tiered'):
            self.assertEqual(sizeAfterCompact, sizeAfterNewInserts)

if __name__ == '__main__':
    wttest.run()


""" Code for computing SW distances between PDs [1]_ of point cloud summaries of activations

Notes
-----
Relevant section : Experiments with PH

Relevant library : `Persim` [2]_

References
----------
.. [1] Carrière, M.; Cuturi, M.; and Oudot, S. 2017. Sliced Wasserstein Kernel for
   Persistence Diagrams. In Precup, D.; and Teh, Y. W., eds., Proceedings of the
   34th International Conference on Machine Learning, volume 70 of Proceedings of
   Machine Learning Research, 664–673. PMLR.
.. [2] Saul, N.; and Tralie, C. 2019. Scikit-TDA: Topological Data Analysis for Python.
"""

import pickle
import argparse
import numpy as np

from persim import sliced_wasserstein

# UPDATE THESE TO REFLECT YOUR OWN DIRECTORIES AND FILE NAMING CONVENTIONS:
# path to project directory containing all model and experiment files
projectdir = '/rcfs/projects/blaktop'
# path to experiment directory containing PH results
expdir = f'{projectdir}/resnet_cifar_experiments'
# directory prefix and filename suffix for PH results per model/batch
# e.g., path to PH results for model i on batch b is expdir/prefix_{i}/persistence_batch{b}filesuffix.p
prefix = 'resnet18_cifar_large'
filesuffix = '_1000'

# number of randomly initialized models (used in 'cross' mode, see below)
num_models = 100

def get_layers(PH):
    """ Returns layers from PH dict keys in the correct order

    Note
    ----
    Specifically designed for the module names defined in `cifar_resnet.resnet18`
    """
    # use key so that conv1 is first (before all block_seq)
    return sorted(PH, key = lambda x : x.replace('conv',''))

def SW_dist_internal(PH, layers):
    """ Computes SW distance between layers of a model """
    nlayers = len(layers)
    dist = np.zeros((nlayers, nlayers))
    for i, layer_i in enumerate(layers[:-1]):
        for j, layer_j in enumerate(layers[i+1:], start=i+1):
            dist[i][j] = sliced_wasserstein(PH[layer_i]['dgms'][1], PH[layer_j]['dgms'][1])
            dist[j][i] = dist[i][j]
    return dist

def SW_dist_cross_model(PH_i, PH_j, layers):
    """ Computes SW distances between layers for two differently initialized models """
    nlayers = len(layers)
    dist = np.zeros((nlayers, nlayers))
    for i, layer_i in enumerate(layers):
        for j, layer_j in enumerate(layers):
            dist[i][j] = sliced_wasserstein(PH_i[layer_i]['dgms'][1], PH_j[layer_j]['dgms'][1])
    return dist

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='SW Distances between PDs of Point Cloud Summaries')
    parser.add_argument('-fs', type=int, help='run index of first random seed ResNet-18 model', default=0)
    parser.add_argument('-ls', type=int, help='run index of last random seed ResNet-18 model (exclusive)', default=1)
    parser.add_argument('-fb', type=int, help='index of first batch', default=0)
    parser.add_argument('-lb', type=int, help='index of last batch (exclusive)', default=1)
    parser.add_argument('-m', type=str, help="mode: either 'int' for single model internal distances, or 'cross' for distances between differently initialized models", default='int')
    args = parser.parse_args()

    for b in range(args.fb, args.lb):
        filename = f'persistence_batch{b}{filesuffix}.p'
        for i in range(args.fs, args.ls):
            savepath = f'{expdir}/{prefix}_{i}'

            PH = pickle.load(open(f'{savepath}/{filename}','rb'))
            layers = get_layers(PH)

            if args.m == 'int':
                dist = SW_dist_internal(PH, layers)
            else:
                dist = []
                for j in range(num_models):
                    if j != i:
                        otherpath = f'{expdir}/{prefix}_{j}'
                        PH_other = pickle.load(open(f'{otherpath}/{filename}','rb'))
                        dist_other = SW_dist_cross_model(PH, PH_other, layers)
                    else:
                        dist_other = SW_dist_internal(PH,layers)
                    dist.append(dist_other)
                dist = np.concatenate(dist, axis=1)

            np.save(f'{savepath}/sliced_wasserstein_batch{b}{filesuffix}_{args.m}', dist)


'''
函数说明:
Author: hongqing
Date: 2021-08-04 14:23:54
LastEditTime: 2021-08-04 15:23:25
'''
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
numoffinger=21
class Net(nn.Module):
    def __init__(self,type=1):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(numoffinger-1, 255)
        self.fc2 = nn.Linear(255, 255)
        self.fc3 = nn.Linear(255, 3)
        if(type==1):
            self.fc1.weight.data.normal_(0, 3)   # initialization
            self.fc2.weight.data.normal_(0, 3)
            self.fc3.weight.data.normal_(0, 3)
        if(type==0):
            self.fc1.weight.data.zero_()   # initialization
            self.fc2.weight.data.zero_()
            self.fc3.weight.data.zero_()
        if(type==2):
            self.fc1.weight.data.random_(1,2)   # initialization
            self.fc2.weight.data.random_(1,2)
            self.fc3.weight.data.random_(1,2)
    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        x = F.relu(x)
        x = self.fc3(x)
        return x


"""empty message

Revision ID: c4665b8d682b
Revises: 10dbb0e0a903
Create Date: 2019-12-26 14:38:11.609539

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = 'c4665b8d682b'
down_revision = '10dbb0e0a903'
branch_labels = None
depends_on = None

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('goods', sa.Column('gage', sa.Integer(), nullable=True))
    # ### end Alembic commands ###

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('goods', 'gage')
    # ### end Alembic commands ###


import requests
import json
import urllib3
from settings import settings as settings
from sdwan_operations import monitor as sdwanmn
import time
import sys, getopt

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

def main(argv):
    # To run the program use the syntax:
    # python __main__.py -d <devfile> -f <fieldfile> -c <optional csvfile>
    devfile = ''
    fieldfile = ''
    csvfile = ''
    try:
        opts, args = getopt.getopt(argv,"hd:f:c:",["devfile=", "fieldfile=", "csvfile"])
    except getopt.GetoptError:
        print ('python __main__.py -d <devfile> -f <fieldfile> -c <optional csvfile>')
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-h':
            print('python __main__.py -d <devfile> -f <fieldfile> -c <optional csvfile>')
            sys.exit(1)
        elif opt in ("-d", "--devfile"):
            devfile = arg
        elif opt in ("-f", "--fieldfile"):
            fieldfile = arg
        elif opt in ("-c", "--csvfile"):
            csvfile = arg

    # Read the device list from the supplied file (for example, dev_list.json)
    try:
        with open(devfile) as f:
            dev_list = json.load(f)
    except FileNotFoundError:
        print ('python __main__.py -d <devfile> -f <fieldfile> -c <optional csvfile>')
        sys.exit(2)

    # Read the field list from the supplied file (for example, field_list.json)
    try:
        with open(fieldfile) as f:
            field_list = json.load(f)
    except FileNotFoundError:
        print ('python __main__.py -d <devfile> -f <fieldfile> -c <optional csvfile>')
        sys.exit(2)
    while (True):
        # Create a session and pass it as a parameter for show_port_stats
        session = requests.Session()
        session.auth = (settings.vmanage_username, settings.vmanage_password)
        sdwanmn.show_port_stats(session, dev_list, field_list)
        # Comment below to stop exporting data to csv
        if csvfile != "":
            sdwanmn.dump_to_csv(session,dev_list,field_list,csvfile)
        time.sleep(60)

if __name__ == "__main__":
    main(sys.argv[1:])


"""
The codes are heavily borrowed from NeuS
"""

import os
import cv2 as cv
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import logging
import mcubes
from icecream import ic
from models.render_utils import sample_pdf

from models.projector import Projector
from tsparse.torchsparse_utils import sparse_to_dense_channel

from models.fast_renderer import FastRenderer

from models.patch_projector import PatchProjector

class SparseNeuSRenderer(nn.Module):
    """
    conditional neus render;
    optimize on normalized world space;
    warped by nn.Module to support DataParallel traning
    """

    def __init__(self,
                 rendering_network_outside,
                 sdf_network,
                 variance_network,
                 rendering_network,
                 n_samples,
                 n_importance,
                 n_outside,
                 perturb,
                 alpha_type='div',
                 conf=None
                 ):
        super(SparseNeuSRenderer, self).__init__()

        self.conf = conf
        self.base_exp_dir = conf['general.base_exp_dir']

        # network setups
        self.rendering_network_outside = rendering_network_outside
        self.sdf_network = sdf_network
        self.variance_network = variance_network
        self.rendering_network = rendering_network

        self.n_samples = n_samples
        self.n_importance = n_importance
        self.n_outside = n_outside
        self.perturb = perturb
        self.alpha_type = alpha_type

        self.rendering_projector = Projector()  # used to obtain features for generalized rendering

        self.h_patch_size = self.conf.get_int('model.h_patch_size', default=3)
        self.patch_projector = PatchProjector(self.h_patch_size)

        self.ray_tracer = FastRenderer()  # ray_tracer to extract depth maps from sdf_volume

        # - fitted rendering or general rendering
        try:
            self.if_fitted_rendering = self.sdf_network.if_fitted_rendering
        except:
            self.if_fitted_rendering = False

    def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_variance,
                  conditional_valid_mask_volume=None):
        device = rays_o.device
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3

        if conditional_valid_mask_volume is not None:
            pts_mask = self.get_pts_mask_for_conditional_volume(pts.view(-1, 3), conditional_valid_mask_volume)
            pts_mask = pts_mask.reshape(batch_size, n_samples)
            pts_mask = pts_mask[:, :-1] * pts_mask[:, 1:]  # [batch_size, n_samples-1]
        else:
            pts_mask = torch.ones([batch_size, n_samples]).to(pts.device)

        sdf = sdf.reshape(batch_size, n_samples)
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
        mid_sdf = (prev_sdf + next_sdf) * 0.5
        dot_val = None
        if self.alpha_type == 'uniform':
            dot_val = torch.ones([batch_size, n_samples - 1]) * -1.0
        else:
            dot_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)
            prev_dot_val = torch.cat([torch.zeros([batch_size, 1]).to(device), dot_val[:, :-1]], dim=-1)
            dot_val = torch.stack([prev_dot_val, dot_val], dim=-1)
            dot_val, _ = torch.min(dot_val, dim=-1, keepdim=False)
            dot_val = dot_val.clip(-10.0, 0.0) * pts_mask
        dist = (next_z_vals - prev_z_vals)
        prev_esti_sdf = mid_sdf - dot_val * dist * 0.5
        next_esti_sdf = mid_sdf + dot_val * dist * 0.5
        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_variance)
        next_cdf = torch.sigmoid(next_esti_sdf * inv_variance)
        alpha_sdf = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)

        alpha = alpha_sdf

        # - apply pts_mask
        alpha = pts_mask * alpha

        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1]).to(device), 1. - alpha + 1e-7], -1), -1)[:, :-1]

        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
        return z_samples

    def cat_z_vals(self, rays_o, rays_d, z_vals, new_z_vals, sdf, lod,
                   sdf_network, gru_fusion,
                   # * related to conditional feature
                   conditional_volume=None,
                   conditional_valid_mask_volume=None
                   ):
        device = rays_o.device
        batch_size, n_samples = z_vals.shape
        _, n_importance = new_z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * new_z_vals[..., :, None]

        if conditional_valid_mask_volume is not None:
            pts_mask = self.get_pts_mask_for_conditional_volume(pts.view(-1, 3), conditional_valid_mask_volume)
            pts_mask = pts_mask.reshape(batch_size, n_importance)
            pts_mask_bool = (pts_mask > 0).view(-1)
        else:
            pts_mask = torch.ones([batch_size, n_importance]).to(pts.device)

        new_sdf = torch.ones([batch_size * n_importance, 1]).to(pts.dtype).to(device) * 100

        if torch.sum(pts_mask) > 1:
            new_outputs = sdf_network.sdf(pts.reshape(-1, 3)[pts_mask_bool], conditional_volume, lod=lod)
            new_sdf[pts_mask_bool] = new_outputs['sdf_pts_scale%d' % lod]  # .reshape(batch_size, n_importance)

        new_sdf = new_sdf.view(batch_size, n_importance)

        z_vals = torch.cat([z_vals, new_z_vals], dim=-1)
        sdf = torch.cat([sdf, new_sdf], dim=-1)

        z_vals, index = torch.sort(z_vals, dim=-1)
        xx = torch.arange(batch_size)[:, None].expand(batch_size, n_samples + n_importance).reshape(-1)
        index = index.reshape(-1)
        sdf = sdf[(xx, index)].reshape(batch_size, n_samples + n_importance)

        return z_vals, sdf

    @torch.no_grad()
    def get_pts_mask_for_conditional_volume(self, pts, mask_volume):
        """

        :param pts: [N, 3]
        :param mask_volume: [1, 1, X, Y, Z]
        :return:
        """
        num_pts = pts.shape[0]
        pts = pts.view(1, 1, 1, num_pts, 3)  # - should be in range (-1, 1)

        pts = torch.flip(pts, dims=[-1])

        pts_mask = F.grid_sample(mask_volume, pts, mode='nearest')  # [1, c, 1, 1, num_pts]
        pts_mask = pts_mask.view(-1, num_pts).permute(1, 0).contiguous()  # [num_pts, 1]

        return pts_mask

    def render_core(self,
                    rays_o,
                    rays_d,
                    z_vals,
                    sample_dist,
                    lod,
                    sdf_network,
                    rendering_network,
                    background_alpha=None,  # - no use here
                    background_sampled_color=None,  # - no use here
                    background_rgb=None,  # - no use here
                    alpha_inter_ratio=0.0,
                    # * related to conditional feature
                    conditional_volume=None,
                    conditional_valid_mask_volume=None,
                    # * 2d feature maps
                    feature_maps=None,
                    color_maps=None,
                    w2cs=None,
                    intrinsics=None,
                    img_wh=None,
                    query_c2w=None,  # - used for testing
                    if_general_rendering=True,
                    if_render_with_grad=True,
                    # * used for blending mlp rendering network
                    img_index=None,
                    rays_uv=None,
                    # * used for clear bg and fg
                    bg_num=0
                    ):
        device = rays_o.device
        N_rays = rays_o.shape[0]
        _, n_samples = z_vals.shape
        dists = z_vals[..., 1:] - z_vals[..., :-1]
        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape).to(device)], -1)

        mid_z_vals = z_vals + dists * 0.5
        mid_dists = mid_z_vals[..., 1:] - mid_z_vals[..., :-1]

        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # n_rays, n_samples, 3
        dirs = rays_d[:, None, :].expand(pts.shape)

        pts = pts.reshape(-1, 3)
        dirs = dirs.reshape(-1, 3)

        # * if conditional_volume is restored from sparse volume, need mask for pts
        if conditional_valid_mask_volume is not None:
            pts_mask = self.get_pts_mask_for_conditional_volume(pts, conditional_valid_mask_volume)
            pts_mask = pts_mask.reshape(N_rays, n_samples).float().detach()
            pts_mask_bool = (pts_mask > 0).view(-1)

            if torch.sum(pts_mask_bool.float()) < 1:  # ! when render out image, may meet this problem
                pts_mask_bool[:100] = True

        else:
            pts_mask = torch.ones([N_rays, n_samples]).to(pts.device)
        # import ipdb; ipdb.set_trace()
        # pts_valid = pts[pts_mask_bool]
        sdf_nn_output = sdf_network.sdf(pts[pts_mask_bool], conditional_volume, lod=lod)

        sdf = torch.ones([N_rays * n_samples, 1]).to(pts.dtype).to(device) * 100
        sdf[pts_mask_bool] = sdf_nn_output['sdf_pts_scale%d' % lod]  # [N_rays*n_samples, 1]
        feature_vector_valid = sdf_nn_output['sdf_features_pts_scale%d' % lod]
        feature_vector = torch.zeros([N_rays * n_samples, feature_vector_valid.shape[1]]).to(pts.dtype).to(device)
        feature_vector[pts_mask_bool] = feature_vector_valid

        # * estimate alpha from sdf
        gradients = torch.zeros([N_rays * n_samples, 3]).to(pts.dtype).to(device)
        # import ipdb; ipdb.set_trace()
        gradients[pts_mask_bool] = sdf_network.gradient(
            pts[pts_mask_bool], conditional_volume, lod=lod).squeeze()

        sampled_color_mlp = None
        rendering_valid_mask_mlp = None
        sampled_color_patch = None
        rendering_patch_mask = None

        if self.if_fitted_rendering:  # used for fine-tuning
            position_latent = sdf_nn_output['sampled_latent_scale%d' % lod]
            sampled_color_mlp = torch.zeros([N_rays * n_samples, 3]).to(pts.dtype).to(device)
            sampled_color_mlp_mask = torch.zeros([N_rays * n_samples, 1]).to(pts.dtype).to(device)

            # - extract pixel
            pts_pixel_color, pts_pixel_mask = self.patch_projector.pixel_warp(
                pts[pts_mask_bool][:, None, :], color_maps, intrinsics,
                w2cs, img_wh=None)  # [N_rays * n_samples,1, N_views,  3] , [N_rays*n_samples, 1, N_views]
            pts_pixel_color = pts_pixel_color[:, 0, :, :]  # [N_rays * n_samples, N_views,  3]
            pts_pixel_mask = pts_pixel_mask[:, 0, :]  # [N_rays*n_samples, N_views]

            # - extract patch
            if_patch_blending = False if rays_uv is None else True
            pts_patch_color, pts_patch_mask = None, None
            if if_patch_blending:
                pts_patch_color, pts_patch_mask = self.patch_projector.patch_warp(
                    pts.reshape([N_rays, n_samples, 3]),
                    rays_uv, gradients.reshape([N_rays, n_samples, 3]),
                    color_maps,
                    intrinsics[0], intrinsics,
                    query_c2w[0], torch.inverse(w2cs), img_wh=None
                )  # (N_rays, n_samples, N_src, Npx, 3), (N_rays, n_samples, N_src, Npx)
                N_src, Npx = pts_patch_mask.shape[2:]
                pts_patch_color = pts_patch_color.view(N_rays * n_samples, N_src, Npx, 3)[pts_mask_bool]
                pts_patch_mask = pts_patch_mask.view(N_rays * n_samples, N_src, Npx)[pts_mask_bool]

                sampled_color_patch = torch.zeros([N_rays * n_samples, Npx, 3]).to(device)
                sampled_color_patch_mask = torch.zeros([N_rays * n_samples, 1]).to(device)

            sampled_color_mlp_, sampled_color_mlp_mask_, \
            sampled_color_patch_, sampled_color_patch_mask_ = sdf_network.color_blend(
                pts[pts_mask_bool],
                position_latent,
                gradients[pts_mask_bool],
                dirs[pts_mask_bool],
                feature_vector[pts_mask_bool],
                img_index=img_index,
                pts_pixel_color=pts_pixel_color,
                pts_pixel_mask=pts_pixel_mask,
                pts_patch_color=pts_patch_color,
                pts_patch_mask=pts_patch_mask

            )  # [n, 3], [n, 1]
            sampled_color_mlp[pts_mask_bool] = sampled_color_mlp_
            sampled_color_mlp_mask[pts_mask_bool] = sampled_color_mlp_mask_.float()
            sampled_color_mlp = sampled_color_mlp.view(N_rays, n_samples, 3)
            sampled_color_mlp_mask = sampled_color_mlp_mask.view(N_rays, n_samples)
            rendering_valid_mask_mlp = torch.mean(pts_mask * sampled_color_mlp_mask, dim=-1, keepdim=True) > 0.5

            # patch blending
            if if_patch_blending:
                sampled_color_patch[pts_mask_bool] = sampled_color_patch_
                sampled_color_patch_mask[pts_mask_bool] = sampled_color_patch_mask_.float()
                sampled_color_patch = sampled_color_patch.view(N_rays, n_samples, Npx, 3)
                sampled_color_patch_mask = sampled_color_patch_mask.view(N_rays, n_samples)
                rendering_patch_mask = torch.mean(pts_mask * sampled_color_patch_mask, dim=-1,
                                                  keepdim=True) > 0.5  # [N_rays, 1]
            else:
                sampled_color_patch, rendering_patch_mask = None, None

        if if_general_rendering:  # used for general training
            # [512, 128, 16]; [4, 512, 128, 59]; [4, 512, 128, 4]
            ren_geo_feats, ren_rgb_feats, ren_ray_diff, ren_mask, _, _ = self.rendering_projector.compute(
                pts.view(N_rays, n_samples, 3),
                # * 3d geometry feature volumes
                geometryVolume=conditional_volume[0],
                geometryVolumeMask=conditional_valid_mask_volume[0],
                # * 2d rendering feature maps
                rendering_feature_maps=feature_maps, # [n_views, 56, 256, 256]
                color_maps=color_maps,
                w2cs=w2cs,
                intrinsics=intrinsics,
                img_wh=img_wh,
                query_img_idx=0,  # the index of the N_views dim for rendering
                query_c2w=query_c2w,
            )

            # (N_rays, n_samples, 3)
            if if_render_with_grad:
                # import ipdb; ipdb.set_trace()
                # [nrays, 3] [nrays, 1]
                sampled_color, rendering_valid_mask = rendering_network(
                    ren_geo_feats, ren_rgb_feats, ren_ray_diff, ren_mask)
                # import ipdb; ipdb.set_trace()
            else:
                with torch.no_grad():
                    sampled_color, rendering_valid_mask = rendering_network(
                        ren_geo_feats, ren_rgb_feats, ren_ray_diff, ren_mask)
        else:
            sampled_color, rendering_valid_mask = None, None

        inv_variance = self.variance_network(feature_vector)[:, :1].clip(1e-6, 1e6)

        true_dot_val = (dirs * gradients).sum(-1, keepdim=True)  # * calculate

        iter_cos = -(F.relu(-true_dot_val * 0.5 + 0.5) * (1.0 - alpha_inter_ratio) + F.relu(
            -true_dot_val) * alpha_inter_ratio)  # always non-positive

        iter_cos = iter_cos * pts_mask.view(-1, 1)

        true_estimate_sdf_half_next = sdf + iter_cos.clip(-10.0, 10.0) * dists.reshape(-1, 1) * 0.5
        true_estimate_sdf_half_prev = sdf - iter_cos.clip(-10.0, 10.0) * dists.reshape(-1, 1) * 0.5

        prev_cdf = torch.sigmoid(true_estimate_sdf_half_prev * inv_variance)
        next_cdf = torch.sigmoid(true_estimate_sdf_half_next * inv_variance)

        p = prev_cdf - next_cdf
        c = prev_cdf

        if self.alpha_type == 'div':
            alpha_sdf = ((p + 1e-5) / (c + 1e-5)).reshape(N_rays, n_samples).clip(0.0, 1.0)
        elif self.alpha_type == 'uniform':
            uniform_estimate_sdf_half_next = sdf - dists.reshape(-1, 1) * 0.5
            uniform_estimate_sdf_half_prev = sdf + dists.reshape(-1, 1) * 0.5
            uniform_prev_cdf = torch.sigmoid(uniform_estimate_sdf_half_prev * inv_variance)
            uniform_next_cdf = torch.sigmoid(uniform_estimate_sdf_half_next * inv_variance)
            uniform_alpha = F.relu(
                (uniform_prev_cdf - uniform_next_cdf + 1e-5) / (uniform_prev_cdf + 1e-5)).reshape(
                N_rays, n_samples).clip(0.0, 1.0)
            alpha_sdf = uniform_alpha
        else:
            assert False

        alpha = alpha_sdf

        # - apply pts_mask
        alpha = alpha * pts_mask

        # pts_radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).reshape(N_rays, n_samples)
        # inside_sphere = (pts_radius < 1.0).float().detach()
        # relax_inside_sphere = (pts_radius < 1.2).float().detach()
        inside_sphere = pts_mask
        relax_inside_sphere = pts_mask

        weights = alpha * torch.cumprod(torch.cat([torch.ones([N_rays, 1]).to(device), 1. - alpha + 1e-7], -1), -1)[:,
                          :-1]  # n_rays, n_samples
        weights_sum = weights.sum(dim=-1, keepdim=True)
        alpha_sum = alpha.sum(dim=-1, keepdim=True)

        if bg_num > 0:
            weights_sum_fg = weights[:, :-bg_num].sum(dim=-1, keepdim=True)
        else:
            weights_sum_fg = weights_sum

        if sampled_color is not None:
            color = (sampled_color * weights[:, :, None]).sum(dim=1)
        else:
            color = None
        # import ipdb; ipdb.set_trace()

        if background_rgb is not None and color is not None:
            color = color + background_rgb * (1.0 - weights_sum)
            # print("color device:" + str(color.device))
        # if color is not None:
        #     # import ipdb; ipdb.set_trace()
        #     color = color + (1.0 - weights_sum)

        ###################*  mlp color rendering  #####################
        color_mlp = None
        # import ipdb; ipdb.set_trace()
        if sampled_color_mlp is not None:
            color_mlp = (sampled_color_mlp * weights[:, :, None]).sum(dim=1)

        if background_rgb is not None and color_mlp is not None:
            color_mlp = color_mlp + background_rgb * (1.0 - weights_sum)

        ############################ *  patch blending  ################
        blended_color_patch = None
        if sampled_color_patch is not None:
            blended_color_patch = (sampled_color_patch * weights[:, :, None, None]).sum(dim=1)  # [N_rays, Npx, 3]

        ######################################################

        gradient_error = (torch.linalg.norm(gradients.reshape(N_rays, n_samples, 3), ord=2,
                                            dim=-1) - 1.0) ** 2
        # ! the gradient normal should be masked out, the pts out of the bounding box should also be penalized
        gradient_error = (pts_mask * gradient_error).sum() / (
                (pts_mask).sum() + 1e-5)

        depth = (mid_z_vals * weights[:, :n_samples]).sum(dim=1, keepdim=True)
        # print("[TEST]: weights_sum in render_core", weights_sum.mean())
        # print("[TEST]: weights_sum in render_core NAN number", weights_sum.isnan().sum())
        # if weights_sum.isnan().sum() > 0:
        #     import ipdb; ipdb.set_trace()
        return {
            'color': color,
            'color_mask': rendering_valid_mask,  # (N_rays, 1)
            'color_mlp': color_mlp,
            'color_mlp_mask': rendering_valid_mask_mlp,
            'sdf': sdf,  # (N_rays, n_samples)
            'depth': depth,  # (N_rays, 1)
            'dists': dists,
            'gradients': gradients.reshape(N_rays, n_samples, 3),
            'variance': 1.0 / inv_variance,
            'mid_z_vals': mid_z_vals,
            'weights': weights,
            'weights_sum': weights_sum,
            'alpha_sum': alpha_sum,
            'alpha_mean': alpha.mean(),
            'cdf': c.reshape(N_rays, n_samples),
            'gradient_error': gradient_error,
            'inside_sphere': inside_sphere,
            'blended_color_patch': blended_color_patch,
            'blended_color_patch_mask': rendering_patch_mask,
            'weights_sum_fg': weights_sum_fg
        }

    def render(self, rays_o, rays_d, near, far, sdf_network, rendering_network,
               perturb_overwrite=-1,
               background_rgb=None,
               alpha_inter_ratio=0.0,
               # * related to conditional feature
               lod=None,
               conditional_volume=None,
               conditional_valid_mask_volume=None,
               # * 2d feature maps
               feature_maps=None,
               color_maps=None,
               w2cs=None,
               intrinsics=None,
               img_wh=None,
               query_c2w=None,  # -used for testing
               if_general_rendering=True,
               if_render_with_grad=True,
               # * used for blending mlp rendering network
               img_index=None,
               rays_uv=None,
               # * importance sample for second lod network
               pre_sample=False,  # no use here
               # * for clear foreground
               bg_ratio=0.0
               ):
        device = rays_o.device
        N_rays = len(rays_o)
        # sample_dist = 2.0 / self.n_samples
        sample_dist = ((far - near) / self.n_samples).mean().item()
        z_vals = torch.linspace(0.0, 1.0, self.n_samples).to(device)
        z_vals = near + (far - near) * z_vals[None, :]

        bg_num = int(self.n_samples * bg_ratio)

        if z_vals.shape[0] == 1:
            z_vals = z_vals.repeat(N_rays, 1)

        if bg_num > 0:
            z_vals_bg = z_vals[:, self.n_samples - bg_num:]
            z_vals = z_vals[:, :self.n_samples - bg_num]

        n_samples = self.n_samples - bg_num
        perturb = self.perturb

        # - significantly speed up training, for the second lod network
        if pre_sample:
            z_vals = self.sample_z_vals_from_maskVolume(rays_o, rays_d, near, far,
                                                        conditional_valid_mask_volume)

        if perturb_overwrite >= 0:
            perturb = perturb_overwrite
        if perturb > 0:
            # get intervals between samples
            mids = .5 * (z_vals[..., 1:] + z_vals[..., :-1])
            upper = torch.cat([mids, z_vals[..., -1:]], -1)
            lower = torch.cat([z_vals[..., :1], mids], -1)
            # stratified samples in those intervals
            t_rand = torch.rand(z_vals.shape).to(device)
            z_vals = lower + (upper - lower) * t_rand

        background_alpha = None
        background_sampled_color = None
        z_val_before = z_vals.clone()
        # Up sample
        if self.n_importance > 0:
            with torch.no_grad():
                pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]

                sdf_outputs = sdf_network.sdf(
                    pts.reshape(-1, 3), conditional_volume, lod=lod)
                # pdb.set_trace()
                sdf = sdf_outputs['sdf_pts_scale%d' % lod].reshape(N_rays, self.n_samples - bg_num)

                n_steps = 4
                for i in range(n_steps):
                    new_z_vals = self.up_sample(rays_o, rays_d, z_vals, sdf, self.n_importance // n_steps,
                                                64 * 2 ** i,
                                                conditional_valid_mask_volume=conditional_valid_mask_volume,
                                                )

                    # if new_z_vals.isnan().sum() > 0:
                    #     import ipdb; ipdb.set_trace()

                    z_vals, sdf = self.cat_z_vals(
                        rays_o, rays_d, z_vals, new_z_vals, sdf, lod,
                        sdf_network, gru_fusion=False,
                        conditional_volume=conditional_volume,
                        conditional_valid_mask_volume=conditional_valid_mask_volume,
                    )

                del sdf

            n_samples = self.n_samples + self.n_importance

        # Background
        ret_outside = None

        # Render
        if bg_num > 0:
            z_vals = torch.cat([z_vals, z_vals_bg], dim=1)
        # if z_vals.isnan().sum() > 0:
        #     import ipdb; ipdb.set_trace()
        ret_fine = self.render_core(rays_o,
                                    rays_d,
                                    z_vals,
                                    sample_dist,
                                    lod,
                                    sdf_network,
                                    rendering_network,
                                    background_rgb=background_rgb,
                                    background_alpha=background_alpha,
                                    background_sampled_color=background_sampled_color,
                                    alpha_inter_ratio=alpha_inter_ratio,
                                    # * related to conditional feature
                                    conditional_volume=conditional_volume,
                                    conditional_valid_mask_volume=conditional_valid_mask_volume,
                                    # * 2d feature maps
                                    feature_maps=feature_maps,
                                    color_maps=color_maps,
                                    w2cs=w2cs,
                                    intrinsics=intrinsics,
                                    img_wh=img_wh,
                                    query_c2w=query_c2w,
                                    if_general_rendering=if_general_rendering,
                                    if_render_with_grad=if_render_with_grad,
                                    # * used for blending mlp rendering network
                                    img_index=img_index,
                                    rays_uv=rays_uv
                                    )

        color_fine = ret_fine['color']

        if self.n_outside > 0:
            color_fine_mask = torch.logical_or(ret_fine['color_mask'], ret_outside['color_mask'])
        else:
            color_fine_mask = ret_fine['color_mask']

        weights = ret_fine['weights']
        weights_sum = ret_fine['weights_sum']

        gradients = ret_fine['gradients']
        mid_z_vals = ret_fine['mid_z_vals']

        # depth = (mid_z_vals * weights[:, :n_samples]).sum(dim=1, keepdim=True)
        depth = ret_fine['depth']
        depth_varaince = ((mid_z_vals - depth) ** 2 * weights[:, :n_samples]).sum(dim=-1, keepdim=True)
        variance = ret_fine['variance'].reshape(N_rays, n_samples).mean(dim=-1, keepdim=True)

        # - randomly sample points from the volume, and maximize the sdf
        pts_random = torch.rand([1024, 3]).float().to(device) * 2 - 1  # normalized to (-1, 1)
        sdf_random = sdf_network.sdf(pts_random, conditional_volume, lod=lod)['sdf_pts_scale%d' % lod]

        result = {
            'depth': depth,
            'color_fine': color_fine,
            'color_fine_mask': color_fine_mask,
            'color_outside': ret_outside['color'] if ret_outside is not None else None,
            'color_outside_mask': ret_outside['color_mask'] if ret_outside is not None else None,
            'color_mlp': ret_fine['color_mlp'],
            'color_mlp_mask': ret_fine['color_mlp_mask'],
            'variance': variance.mean(),
            'cdf_fine': ret_fine['cdf'],
            'depth_variance': depth_varaince,
            'weights_sum': weights_sum,
            'weights_max': torch.max(weights, dim=-1, keepdim=True)[0],
            'alpha_sum': ret_fine['alpha_sum'].mean(),
            'alpha_mean': ret_fine['alpha_mean'],
            'gradients': gradients,
            'weights': weights,
            'gradient_error_fine': ret_fine['gradient_error'],
            'inside_sphere': ret_fine['inside_sphere'],
            'sdf': ret_fine['sdf'],
            'sdf_random': sdf_random,
            'blended_color_patch': ret_fine['blended_color_patch'],
            'blended_color_patch_mask': ret_fine['blended_color_patch_mask'],
            'weights_sum_fg': ret_fine['weights_sum_fg']
        }

        return result

    @torch.no_grad()
    def sample_z_vals_from_sdfVolume(self, rays_o, rays_d, near, far, sdf_volume, mask_volume):
        # ? based on sdf to do importance sampling, seems that too biased on pre-estimation
        device = rays_o.device
        N_rays = len(rays_o)
        n_samples = self.n_samples * 2

        z_vals = torch.linspace(0.0, 1.0, n_samples).to(device)
        z_vals = near + (far - near) * z_vals[None, :]

        if z_vals.shape[0] == 1:
            z_vals = z_vals.repeat(N_rays, 1)

        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]

        sdf = self.get_pts_mask_for_conditional_volume(pts.view(-1, 3), sdf_volume).reshape([N_rays, n_samples])

        new_z_vals = self.up_sample(rays_o, rays_d, z_vals, sdf, self.n_samples,
                                    200,
                                    conditional_valid_mask_volume=mask_volume,
                                    )
        return new_z_vals

    @torch.no_grad()
    def sample_z_vals_from_maskVolume(self, rays_o, rays_d, near, far, mask_volume):  # don't use
        device = rays_o.device
        N_rays = len(rays_o)
        n_samples = self.n_samples * 2

        z_vals = torch.linspace(0.0, 1.0, n_samples).to(device)
        z_vals = near + (far - near) * z_vals[None, :]

        if z_vals.shape[0] == 1:
            z_vals = z_vals.repeat(N_rays, 1)

        mid_z_vals = (z_vals[:, 1:] + z_vals[:, :-1]) * 0.5

        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]

        pts_mask = self.get_pts_mask_for_conditional_volume(pts.view(-1, 3), mask_volume).reshape(
            [N_rays, n_samples - 1])

        # empty voxel set to 0.1, non-empty voxel set to 1
        weights = torch.where(pts_mask > 0, torch.ones_like(pts_mask).to(device),
                              0.1 * torch.ones_like(pts_mask).to(device))

        # sample more pts in non-empty voxels
        z_samples = sample_pdf(z_vals, weights, self.n_samples, det=True).detach()
        return z_samples

    @torch.no_grad()
    def filter_pts_by_depthmaps(self, coords, pred_depth_maps, proj_matrices,
                                partial_vol_origin, voxel_size,
                                near, far, depth_interval, d_plane_nums):
        """
        Use the pred_depthmaps to remove redundant pts (pruned by sdf, sdf always have two sides, the back side is useless)
        :param coords: [n, 3]  int coords
        :param pred_depth_maps: [N_views, 1, h, w]
        :param proj_matrices: [N_views, 4, 4]
        :param partial_vol_origin: [3]
        :param voxel_size: 1
        :param near: 1
        :param far: 1
        :param depth_interval: 1
        :param d_plane_nums: 1
        :return:
        """
        device = pred_depth_maps.device
        n_views, _, sizeH, sizeW = pred_depth_maps.shape

        if len(partial_vol_origin.shape) == 1:
            partial_vol_origin = partial_vol_origin[None, :]
        pts = coords * voxel_size + partial_vol_origin

        rs_grid = pts.unsqueeze(0).expand(n_views, -1, -1)
        rs_grid = rs_grid.permute(0, 2, 1).contiguous()  # [n_views, 3, n_pts]
        nV = rs_grid.shape[-1]
        rs_grid = torch.cat([rs_grid, torch.ones([n_views, 1, nV]).to(device)], dim=1)  # [n_views, 4, n_pts]

        # Project grid
        im_p = proj_matrices @ rs_grid  # - transform world pts to image UV space   # [n_views, 4, n_pts]
        im_x, im_y, im_z = im_p[:, 0], im_p[:, 1], im_p[:, 2]
        im_x = im_x / im_z
        im_y = im_y / im_z

        im_grid = torch.stack([2 * im_x / (sizeW - 1) - 1, 2 * im_y / (sizeH - 1) - 1], dim=-1)

        im_grid = im_grid.view(n_views, 1, -1, 2)
        sampled_depths = torch.nn.functional.grid_sample(pred_depth_maps, im_grid, mode='bilinear',
                                                         padding_mode='zeros',
                                                         align_corners=True)[:, 0, 0, :]  # [n_views, n_pts]
        sampled_depths_valid = (sampled_depths > 0.5 * near).float()
        valid_d_min = (sampled_depths - d_plane_nums * depth_interval).clamp(near.item(),
                                                                             far.item()) * sampled_depths_valid
        valid_d_max = (sampled_depths + d_plane_nums * depth_interval).clamp(near.item(),
                                                                             far.item()) * sampled_depths_valid

        mask = im_grid.abs() <= 1
        mask = mask[:, 0]  # [n_views, n_pts, 2]
        mask = (mask.sum(dim=-1) == 2) & (im_z > valid_d_min) & (im_z < valid_d_max)

        mask = mask.view(n_views, -1)
        mask = mask.permute(1, 0).contiguous()  # [num_pts, nviews]

        mask_final = torch.sum(mask.float(), dim=1, keepdim=False) > 0

        return mask_final

    @torch.no_grad()
    def get_valid_sparse_coords_by_sdf_depthfilter(self, sdf_volume, coords_volume, mask_volume, feature_volume,
                                                   pred_depth_maps, proj_matrices,
                                                   partial_vol_origin, voxel_size,
                                                   near, far, depth_interval, d_plane_nums,
                                                   threshold=0.02, maximum_pts=110000):
        """
        assume batch size == 1, from the first lod to get sparse voxels
        :param sdf_volume: [1, X, Y, Z]
        :param coords_volume: [3, X, Y, Z]
        :param mask_volume: [1, X, Y, Z]
        :param feature_volume: [C, X, Y, Z]
        :param threshold:
        :return:
        """
        device = coords_volume.device
        _, dX, dY, dZ = coords_volume.shape

        def prune(sdf_pts, coords_pts, mask_volume, threshold):
            occupancy_mask = (torch.abs(sdf_pts) < threshold).squeeze(1)  # [num_pts]
            valid_coords = coords_pts[occupancy_mask]

            # - filter backside surface by depth maps
            mask_filtered = self.filter_pts_by_depthmaps(valid_coords, pred_depth_maps, proj_matrices,
                                                         partial_vol_origin, voxel_size,
                                                         near, far, depth_interval, d_plane_nums)
            valid_coords = valid_coords[mask_filtered]

            # - dilate
            occupancy_mask = sparse_to_dense_channel(valid_coords, 1, [dX, dY, dZ], 1, 0, device)  # [dX, dY, dZ, 1]

            # - dilate
            occupancy_mask = occupancy_mask.float()
            occupancy_mask = occupancy_mask.view(1, 1, dX, dY, dZ)
            occupancy_mask = F.avg_pool3d(occupancy_mask, kernel_size=7, stride=1, padding=3)
            occupancy_mask = occupancy_mask.view(-1, 1) > 0

            final_mask = torch.logical_and(mask_volume, occupancy_mask)[:, 0]  # [num_pts]

            return final_mask, torch.sum(final_mask.float())

        C, dX, dY, dZ = feature_volume.shape
        sdf_volume = sdf_volume.permute(1, 2, 3, 0).contiguous().view(-1, 1)
        coords_volume = coords_volume.permute(1, 2, 3, 0).contiguous().view(-1, 3)
        mask_volume = mask_volume.permute(1, 2, 3, 0).contiguous().view(-1, 1)
        feature_volume = feature_volume.permute(1, 2, 3, 0).contiguous().view(-1, C)

        # - for check
        # sdf_volume = torch.rand_like(sdf_volume).float().to(sdf_volume.device) * 0.02

        final_mask, valid_num = prune(sdf_volume, coords_volume, mask_volume, threshold)

        while (valid_num > maximum_pts) and (threshold > 0.003):
            threshold = threshold - 0.002
            final_mask, valid_num = prune(sdf_volume, coords_volume, mask_volume, threshold)

        valid_coords = coords_volume[final_mask]  # [N, 3]
        valid_feature = feature_volume[final_mask]  # [N, C]

        valid_coords = torch.cat([torch.ones([valid_coords.shape[0], 1]).to(valid_coords.device) * 0,
                                  valid_coords], dim=1)  # [N, 4], append batch idx

        # ! if the valid_num is still larger than maximum_pts, sample part of pts
        if valid_num > maximum_pts:
            valid_num = valid_num.long()
            occupancy = torch.ones([valid_num]).to(device) > 0
            choice = np.random.choice(valid_num.cpu().numpy(), valid_num.cpu().numpy() - maximum_pts,
                                      replace=False)
            ind = torch.nonzero(occupancy).to(device)
            occupancy[ind[choice]] = False
            valid_coords = valid_coords[occupancy]
            valid_feature = valid_feature[occupancy]

            print(threshold, "randomly sample to save memory")

        return valid_coords, valid_feature

    @torch.no_grad()
    def get_valid_sparse_coords_by_sdf(self, sdf_volume, coords_volume, mask_volume, feature_volume, threshold=0.02,
                                       maximum_pts=110000):
        """
        assume batch size == 1, from the first lod to get sparse voxels
        :param sdf_volume: [num_pts, 1]
        :param coords_volume: [3, X, Y, Z]
        :param mask_volume: [1, X, Y, Z]
        :param feature_volume: [C, X, Y, Z]
        :param threshold:
        :return:
        """

        def prune(sdf_volume, mask_volume, threshold):
            occupancy_mask = torch.abs(sdf_volume) < threshold  # [num_pts, 1]

            # - dilate
            occupancy_mask = occupancy_mask.float()
            occupancy_mask = occupancy_mask.view(1, 1, dX, dY, dZ)
            occupancy_mask = F.avg_pool3d(occupancy_mask, kernel_size=7, stride=1, padding=3)
            occupancy_mask = occupancy_mask.view(-1, 1) > 0

            final_mask = torch.logical_and(mask_volume, occupancy_mask)[:, 0]  # [num_pts]

            return final_mask, torch.sum(final_mask.float())

        C, dX, dY, dZ = feature_volume.shape
        coords_volume = coords_volume.permute(1, 2, 3, 0).contiguous().view(-1, 3)
        mask_volume = mask_volume.permute(1, 2, 3, 0).contiguous().view(-1, 1)
        feature_volume = feature_volume.permute(1, 2, 3, 0).contiguous().view(-1, C)

        final_mask, valid_num = prune(sdf_volume, mask_volume, threshold)

        while (valid_num > maximum_pts) and (threshold > 0.003):
            threshold = threshold - 0.002
            final_mask, valid_num = prune(sdf_volume, mask_volume, threshold)

        valid_coords = coords_volume[final_mask]  # [N, 3]
        valid_feature = feature_volume[final_mask]  # [N, C]

        valid_coords = torch.cat([torch.ones([valid_coords.shape[0], 1]).to(valid_coords.device) * 0,
                                  valid_coords], dim=1)  # [N, 4], append batch idx

        # ! if the valid_num is still larger than maximum_pts, sample part of pts
        if valid_num > maximum_pts:
            device = sdf_volume.device
            valid_num = valid_num.long()
            occupancy = torch.ones([valid_num]).to(device) > 0
            choice = np.random.choice(valid_num.cpu().numpy(), valid_num.cpu().numpy() - maximum_pts,
                                      replace=False)
            ind = torch.nonzero(occupancy).to(device)
            occupancy[ind[choice]] = False
            valid_coords = valid_coords[occupancy]
            valid_feature = valid_feature[occupancy]

            print(threshold, "randomly sample to save memory")

        return valid_coords, valid_feature

    @torch.no_grad()
    def extract_fields(self, bound_min, bound_max, resolution, query_func, device,
                       # * related to conditional feature
                       **kwargs
                       ):
        N = 64
        X = torch.linspace(bound_min[0], bound_max[0], resolution).to(device).split(N)
        Y = torch.linspace(bound_min[1], bound_max[1], resolution).to(device).split(N)
        Z = torch.linspace(bound_min[2], bound_max[2], resolution).to(device).split(N)

        u = np.zeros([resolution, resolution, resolution], dtype=np.float32)
        with torch.no_grad():
            for xi, xs in enumerate(X):
                for yi, ys in enumerate(Y):
                    for zi, zs in enumerate(Z):
                        xx, yy, zz = torch.meshgrid(xs, ys, zs, indexing="ij")
                        pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1)

                        # ! attention, the query function is different for extract geometry and fields
                        output = query_func(pts, **kwargs)
                        sdf = output['sdf_pts_scale%d' % kwargs['lod']].reshape(len(xs), len(ys),
                                                                                len(zs)).detach().cpu().numpy()

                        u[xi * N: xi * N + len(xs), yi * N: yi * N + len(ys), zi * N: zi * N + len(zs)] = -1 * sdf
        return u

    @torch.no_grad()
    def extract_geometry(self, sdf_network, bound_min, bound_max, resolution, threshold, device, occupancy_mask=None,
                         # * 3d feature volume
                         **kwargs
                         ):
        # logging.info('threshold: {}'.format(threshold))

        u = self.extract_fields(bound_min, bound_max, resolution,
                                lambda pts, **kwargs: sdf_network.sdf(pts, **kwargs),
                                # - sdf need to be multiplied by -1
                                device,
                                # * 3d feature volume
                                **kwargs
                                )
        if occupancy_mask is not None:
            dX, dY, dZ = occupancy_mask.shape
            empty_mask = 1 - occupancy_mask
            empty_mask = empty_mask.view(1, 1, dX, dY, dZ)
            # - dilation
            # empty_mask = F.avg_pool3d(empty_mask, kernel_size=7, stride=1, padding=3)
            empty_mask = F.interpolate(empty_mask, [resolution, resolution, resolution], mode='nearest')
            empty_mask = empty_mask.view(resolution, resolution, resolution).cpu().numpy() > 0
            u[empty_mask] = -100
            del empty_mask

        vertices, triangles = mcubes.marching_cubes(u, threshold)
        b_max_np = bound_max.detach().cpu().numpy()
        b_min_np = bound_min.detach().cpu().numpy()

        vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
        return vertices, triangles, u

    @torch.no_grad()
    def extract_depth_maps(self, sdf_network, con_volume, intrinsics, c2ws, H, W, near, far):
        """
        extract depth maps from the density volume
        :param con_volume: [1, 1+C, dX, dY, dZ]  can by con_volume or sdf_volume
        :param c2ws: [B, 4, 4]
        :param H:
        :param W:
        :param near:
        :param far:
        :return:
        """
        device = con_volume.device
        batch_size = intrinsics.shape[0]

        with torch.no_grad():
            ys, xs = torch.meshgrid(torch.linspace(0, H - 1, H),
                                    torch.linspace(0, W - 1, W), indexing="ij")  # pytorch's meshgrid has indexing='ij'
            p = torch.stack([xs, ys, torch.ones_like(ys)], dim=-1)  # H, W, 3

            intrinsics_inv = torch.inverse(intrinsics)

            p = p.view(-1, 3).float().to(device)  # N_rays, 3
            p = torch.matmul(intrinsics_inv[:, None, :3, :3], p[:, :, None]).squeeze()  # Batch, N_rays, 3
            rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # Batch, N_rays, 3
            rays_v = torch.matmul(c2ws[:, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # Batch, N_rays, 3
            rays_o = c2ws[:, None, :3, 3].expand(rays_v.shape)  # Batch, N_rays, 3
            rays_d = rays_v

        rays_o = rays_o.contiguous().view(-1, 3)
        rays_d = rays_d.contiguous().view(-1, 3)

        ################## - sphere tracer to extract depth maps               ######################
        depth_masks_sphere, depth_maps_sphere = self.ray_tracer.extract_depth_maps(
            rays_o, rays_d,
            near[None, :].repeat(rays_o.shape[0], 1),
            far[None, :].repeat(rays_o.shape[0], 1),
            sdf_network, con_volume
        )

        depth_maps = depth_maps_sphere.view(batch_size, 1, H, W)
        depth_masks = depth_masks_sphere.view(batch_size, 1, H, W)

        depth_maps = torch.where(depth_masks, depth_maps,
                                 torch.zeros_like(depth_masks.float()).to(device))  # fill invalid pixels by 0

        return depth_maps, depth_masks


from django.db import models

# Create your models here.

class Meal(models.Model):
    menu = models.TextField(blank=True)

    def __str__(self):
        return self.menu

class MealList(models.Model):
    date = models.DateField(blank=True)

    breakfast = models.ForeignKey(
        Meal,
        blank=True,
        null=True,
        on_delete=models.CASCADE,
        related_name="breakfast_id",
    )
    lunch = models.ForeignKey(
        Meal,
        null=True,
        blank=True,
        on_delete=models.CASCADE,
        related_name="lunch"
    )
    dinner = models.ForeignKey(
        Meal,
        null=True,
        blank=True,
        on_delete=models.CASCADE,
        related_name="dinner"
    )

    created_at = models.DateTimeField(auto_now_add = True)
    updated_at = models.DateTimeField(auto_now = True)

    class Meta:
        db_table = 'meal_lists'

    def update_meal(self, data, number):
        if number == 1:
            self.breakfast = data
        if number == 2:
            self.lunch = data
        if number == 3:
            self.dinner = data
        self.save()
        return self


def is_prime(num):
    if(num>1):
        for i in range(2,num):
            if((num % i) != 0):
                #number is prime
                return 1
            else:
                return 0

#driver code

num= int(input("\n Enter the positive number: "))
print(is_prime(num))
if(num==1):
    print("\n 1 is neither prime nor composite. ")
else:
    if(is_prime(num)==1):
        print("\n NUMBER IS PRIME.")
    else:
        print("\n NUMBER IS NOT PRIME. ")


#!/usr/bin/env python
"""Notepad-- : A very simple text editor."""

import sys
if sys.version_info < (3,):
    import Tkinter as tk
    import tkFileDialog as filedialog
else:
    import tkinter as tk
    from tkinter import filedialog

def create_gui(root):
    """Set up the GUI and functionality."""
    root.title("Notepad--")
    textbox = tk.Text(root)
    textbox.pack(expand=True, fill=tk.BOTH)

    # Functionality for opening and saving files.
    def open_file():
        """Open a file"""
        f = filedialog.askopenfile()
        if f:
            textbox.delete(1.0, tk.END)
            textbox.insert(tk.END, f.read())
            f.close()

    def save_file():
        """Save a file"""
        f = filedialog.asksaveasfile()
        if f:
            f.write(textbox.get(1.0, tk.END)[:-1])
            f.close()

    # Set up the command buttons
    frame = tk.Frame(root)
    frame.pack()
    tk.Button(frame, text="Open", command=open_file).pack(side=tk.LEFT)
    tk.Button(frame, text="Save", command=save_file).pack(side=tk.LEFT)

if __name__ == "__main__":
    root = tk.Tk()
    create_gui(root)
    root.mainloop()


#!/usr/bin/python
######################################################################################
## Find disk usage under a root directory, examples:
##   1. python du_rootdir.py -d rootDir
##      Find disk usage under the root directory "rootDir"
##   2. python du_rootdir.py -d rootDir -r true
##      Find disk usage recursively under the root directory "rootDir"
##   3. python du_rootdir.py -d rootDir -r true -l 3
##      Find disk usage recursively under the root directory "rootDir" with
##      recursive level 3
## Author: Zhichang Guo, email: Zhichang.Guo@noaa.gov
######################################################################################
from pathlib import Path
import argparse
import os
import sys

def print_indented(result, level):
    print('\t' * level + result)

def traverse_dir(dir):
    iflag = 'not'
    l = os.listdir(dir)
    for d in l:
        if os.path.isdir(dir + d):
            if not d.startswith('.') or (d.startswith('.') and not iflag.upper() == 'NOT'):
                fullpath = os.path.join(dir, d)
                cmd = "du -sh " + " \'" + str(fullpath) + "\'"
#               print(cmd)
                result = os.popen(cmd).read()
#               print(result)
                print(result,end='')

def traverse_dir_recur(dir, max_level, level=0):
    iflag = 'not'
    l = os.listdir(dir)
    for d in l:
        if os.path.isdir(dir + d):
            if not d.startswith('.') or (d.startswith('.') and not iflag.upper() == 'NOT'):
                fullpath = os.path.join(dir, d)
                if max_level == 0 or level < max_level:
                    traverse_dir_recur(dir + d + "/", max_level, level+1)
                    cmd = "du -sh " + " \'" + str(fullpath) + "\'"
#                   print(cmd)
                    result = os.popen(cmd).read()
#                   print_indented(result.rstrip("\n"), level+1)
                    if level > 0:
                        print('\033[1;31;43m ' + result.rstrip("\n") + ' \033[0;0m')
                    else:
                        print(result.rstrip("\n"))
                    #print(result)

def find_du(rootdir, rflag, max_level):
    home = os.environ['HOME']
    path = Path(rootdir)
    owner = path.owner()
    print("Owner: ",owner)
    rootdir += '/'
    if rflag == 'false':
        traverse_dir(rootdir)
    else:
        traverse_dir_recur(rootdir, max_level)

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument('-d', '--rootd', help="name of the root directory", required=True)
    ap.add_argument('-r', '--rflag', help="recurively or not", default="false")
    ap.add_argument('-l', '--level', help="level", type=int, default=0)
    MyArgs = ap.parse_args()
    find_du(MyArgs.rootd, MyArgs.rflag, MyArgs.level)


from os import path
from subprocess import check_output, check_call
from pathlib import Path

def _filename(fname):
    import yapl
    yapl_root = Path(yapl.__file__).parent
    filename = fname + ".sh"
    full_filename = yapl_root / filename
    return str(full_filename)

def func(fname, *args):
    shell_script = _filename(fname)
    command = list([shell_script] + [str(a) for a in args])
    result = check_output(command, text=True).strip()
    return result

def proc(fname, *args):
    shell_script = _filename(fname)
    command = list([shell_script] + [str(a) for a in args])
    check_call(command)


import os, random
import numpy as np
import torch
import argparse

from train import train

def init_seeds(seed=0):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

def parsing_args(c):
    parser = argparse.ArgumentParser(description='msflow')
    parser.add_argument('--mode', default='train', type=str,
                        help='train or test.')
    parser.add_argument('--resume', action='store_true', default=False,
                        help='resume training or not.')
    parser.add_argument('--eval_ckpt', default='', type=str,
                        help='checkpoint path for evaluation.')
    parser.add_argument('--gpu', default='0', type=str)
    parser.add_argument('--class-name', default='bottle', type=str)
    parser.add_argument('--lr', default=1e-4, type=float,
                        help='learning rate')
    parser.add_argument('--batch-size', default=8, type=int,
                        help='train batch size')
    parser.add_argument('--meta-epochs', default=25, type=int,
                        help='number of meta epochs to train')
    parser.add_argument('--sub-epochs', default=4, type=int,
                        help='number of sub epochs to train')
    parser.add_argument('--extractor', default='wide_resnet50_2', type=str,
                        help='feature extractor')
    parser.add_argument('--pool-type', default='avg', type=str,
                        help='pool type for extracted feature maps')
    parser.add_argument('--parallel-blocks', default=[2, 5, 8], type=int, metavar='L', nargs='+',
                        help='number of flow blocks used in parallel flows.')
    parser.add_argument('--pro-eval', action='store_true', default=False,
                        help='evaluate the pro score or not.')
    parser.add_argument('--pro-eval-interval', default=4, type=int,
                        help='interval for pro evaluation.')

    args = parser.parse_args()

    for k, v in vars(args).items():
        setattr(c, k, v)

    c.input_size = (256, 256) if c.class_name == 'transistor' else (512, 512)

    return c

def main(c):
    c = parsing_args(c)
    os.environ['CUDA_VISIBLE_DEVICES'] = c.gpu
    init_seeds(seed=c.seed)
    c.version_name = 'msflow_{}_{}pool_pl{}'.format(c.extractor, c.pool_type, "".join([str(x) for x in c.parallel_blocks]))
    c.ckpt_dir = os.path.join(c.work_dir, c.version_name, c.class_name)
    train(c)

if __name__ == '__main__':
    import default as c
    main(c)


import os, sys

from langchain.llms import OpenAI
from langchain.chains.question_answering import load_qa_chain
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS

prefix, query = sys.argv[1:3]

api_key = os.environ.get('OPENAI_API_KEY')
if api_key is None:
    sys.exit("OPENAI_API_KEY is unset")

embeddings = OpenAIEmbeddings(openai_api_key=api_key)
llm = OpenAI(temperature=0.7, openai_api_key=api_key)
qa_chain = load_qa_chain(llm, chain_type="stuff")

fdir = os.path.dirname(prefix)
fname = os.path.basename(prefix)
faiss_index = FAISS.load_local(fdir, embeddings, fname)

# This model's maximum context length is 4097 tokens, 256 for the completion
ss = faiss_index.similarity_search(query.strip(), k=5)
ans = qa_chain.run(input_documents=ss, question=query)

print(ans.strip())


print('Задача 1. Результаты')

import os
import random

# Одному программисту дали задачу для обработки неких результатов тестирования
# двух групп людей. Файл первой группы (group_1.txt) находится в папке task,
# файл второй группы (group_2.txt) — в папке Additional_info.

# На экран нужно было вывести сумму очков первой группы, затем разность очков
# опять же первой группы и напоследок — произведение очков уже второй группы.

# Программист оказался не очень опытным, писал код наобум и даже не стал его
# проверять. И оказалось, этот код просто не работает.

file = open('C:\\learn_work\\task\\group_1.txt', 'r', encoding='utf-8')
summa = 0
diff = 0
for i_line in file:
    info = i_line.split()
    summa += int(info[2])
    diff -= int(info[2])
print('Сумма очков первой группы:', summa)
print('Разность очков первой группы:', diff)

file.close()

file_2 = open('C:\\learn_work\\task\\Additional_info\\group_2.txt', 'r', encoding='utf-8')
compose = 1
for i_line in file_2:
    info = i_line.split()
    compose *= int(info[2])

print('Произведение очков второй группы:', compose)
file_2.close()


from turtle import Turtle
import random

COLORS = ["red", "orange", "yellow", "green", "blue", "purple"]
BORDER = "black"
STARTING_MOVE_DISTANCE = 5
MOVE_INCREMENT = 5

class CarManager:
    def __init__(self):
        self.all_cars = []
        self.initial_speed = STARTING_MOVE_DISTANCE

    def createCar(self):
        chance = random.randint(1, 4)  # To create the cars less frequently so that there is space for turtle movement
        if chance == 1:
            new_car = Turtle()
            new_car.speed("fastest")
            new_car.penup()
            new_car.shape("square")
            new_car.shapesize(stretch_wid=1, stretch_len=2)
            new_car.color(BORDER, random.choice(COLORS))
            random_y = random.randint(-230, 230)
            new_car.goto(300, random_y)
            self.all_cars.append(new_car)

    def moveCars(self):
        for car in self.all_cars:
            car.backward(self.initial_speed)

    def fasterCars(self):
        for car in self.all_cars:
            # This is to delete all the cars from a certain level after level is passed
            car.hideturtle()
            car.goto(-350, car.ycor())
            del car
        self.initial_speed += MOVE_INCREMENT


# Пример кода 1. Файл test_add_house

#Асбстрактный сервис по созданию сущности дома и квартир в нем (для ТСЖ)

from model.new_house import New_house
from model.rooms import Rooms
import time

def test_add_house(app):
    app.new_house.fill_info(New_house(district="Железнодорожный", street="Тестовая", house_number="1",
                                      postal_index="630132", tech_status="90",
                                      maintenance_date="01.01.1992", build_date="01.01.1991",
                                      surrounding_area="500", roof_area="250", attic_area="250",
                                      cellar_area="250", stairs_area="300"))

    time.sleep(3)

    app.rooms.fill_flats(Rooms(entrance="1", room_number="1", floor_number="1",
                               total_area="45", living_area="28"))

# Пример кода 2. Файл edit_house

from model.new_house import New_house

def test_modify_house(app):
    app.new_house.modify_first_house(New_house(district="Ленинский", street="ТестоваяИзмененная", house_number="2",
                                     postal_index="630054", tech_status="95", surrounding_area="600",
                                     roof_area="350", attic_area="350", cellar_area="350", stairs_area="200"))


import torch
from torch.optim.optimizer import Optimizer

class COCOB(Optimizer):
    r"""Implements COCOB algorithm.
    It has been proposed in `Training Deep Networks without Learning Rates Through Coin Betting`_.
    Arguments:
        params (iterable): iterable of parameters to optimize or dicts defining
            parameter groups
        alpha (float, optional): It was proposed to increase the stability in the first iterations,
            similarly and independently to the learning rate warm-up. The number roughly denotes the
            number of rounds of warm-up (default 100)
        eps (float, optional): term added to the denominator to improve
            numerical stability (default: 1e-8)
        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
    .. _Training Deep Networks without Learning Rates Through Coin Betting:
        https://arxiv.org/abs/1705.07795
    """

    def __init__(self, params, alpha: float = 100, eps: float = 1e-8, weight_decay: float = 0):
        if not 0.0 <= eps:
            raise ValueError("Invalid epsilon value: {}".format(eps))
        if not 0.0 <= alpha:
            raise ValueError("Invalid alpha value: {}".format(alpha))
        if not 0.0 <= weight_decay:
            raise ValueError(
                "Invalid weight_decay value: {}".format(weight_decay))

        defaults = dict(weight_decay=weight_decay)
        self._alpha = alpha
        self._eps = eps

        super(COCOB, self).__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure = None):
        """Performs a single optimization step.
        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad
                if grad.is_sparse:
                    raise RuntimeError('COCOB does not support sparse gradients')

                state = self.state[p]

                # State initialization
                if len(state) == 0:
                    # Sum of the negative gradients
                    state['sum_negative_gradients'] = torch.zeros_like(p).detach()
                    # Sum of the absolute values of the stochastic subgradients
                    state['grad_norm_sum'] = torch.zeros_like(p).detach()
                    # Maximum observed scale
                    state['L'] = self._eps*torch.ones_like(p).detach()
                    # Reward/wealth of the algorithm for each coordinate
                    state['reward'] = torch.zeros_like(p).detach()
                    # We need to save the initial point because this is a FTRL-based algorithm
                    state['x0'] = torch.clone(p.data).detach()

                sum_negative_gradients, grad_norm_sum, L, reward, x0 = (
                    state['sum_negative_gradients'],
                    state['grad_norm_sum'],
                    state['L'],
                    state['reward'],
                    state['x0'],
                )

                if group['weight_decay'] != 0:
                    grad = grad.add(p, alpha=group['weight_decay'])

                # update maximum rage of the gradients
                torch.max(L, torch.abs(grad), out=L)
                # udpate dual vector
                sum_negative_gradients.sub_(grad)
                # update sum of the absolute values of the gradients
                grad_norm_sum.add_(torch.abs(grad))
                # update the wealth
                reward.addcmul_(grad, p.data.sub(x0), value=-1)
                # reset the wealth to zero in case we lost all
                torch.maximum(reward, torch.zeros_like(reward), out=reward)
                # calculate denominator
                den = torch.maximum(grad_norm_sum.add(L), L.mul(self._alpha)).mul(L)
                # update model parameters
                p.data.copy_(reward.add(L).mul(sum_negative_gradients).div(den).add(x0))

        return loss


import serial
from sense_hat import SenseHat
from socket import gethostname
from xbee import XBee
from statistics import median

def clear_matrix():
    sense.clear()

def show_hostname():
    hostname = gethostname()
    sense.show_message("Hostname: " + hostname)

def receive_data(data):
    print("received data: ", data)
    if (data["rf_data"] == b"ping"):
        # Received ping? -> Send pong back
        send_data("pong")
    elif (data["rf_data"] == b"pong"):
        # Received pong? -> Store & calc RSSI
        rssi_list.append(ord(data["rssi"]))
        print_current_rssi_median()

        if data["source_addr"] not in rssi_list:
            rssi_dict[data["source_addr"]] = []
        rssi_dict[data["source_addr"]].append(ord(data["rssi"]))

def send_data(data, dest_addr="\x00\x0A"):
    xbee.send("tx",
              frame_id="\x00",
              dest_addr=dest_addr,
              data=data)

def init_rssi_calc(n_pings=10):
    rssi_list = []
    dest_addr = "\x00\x0A" # 2byte hex value (TODO: Set according to adress of destination XBee module)
    for i in range(n_pings):
        send_data("ping", dest_addr)

def print_current_rssi_median():
    med = median(rssi_list)
    print("Current RSSI median with {} received pongs: {}".format(len(rssi_list), med))
    dist = dist_from_rssi(med)
    print("Current dist from RSSI: dist = {}".format(dist))

def dist_from_rssi(rssi):
    n = 2.8 # path loss variable from 2 to 4
    A = 33 # TODO: Measure reference RSSI (1m distance)
    # RSSI = -10 * n * log_10(d) + A
    # => Transformed to
    # d = 10^(A/10n)
    dist = 10**(A/10*n)
    return dist

def three_anchor_bbox():
    dist_dict = {}
    for anchor in rssi_dict:
        dist_dict[anchor] = dist_from_rssi(median(rssi_dict[anchor]))

    x = 1/2 * (min([anchor_positions[anchor][0] + dist_dict[anchor] for anchor in anchor_positions]) + \
        max([anchor_positions[anchor][0] - dist_dict[anchor] for anchor in anchor_positions]))

    y = 1/2 * (min([anchor_positions[anchor][1] + dist_dict[anchor] for anchor in anchor_positions]) + \
        max([anchor_positions[anchor][1] - dist_dict[anchor] for anchor in anchor_positions]))

    print(">> BBox calculation done: X: {} | Y: {}".format(x,y))

def three_anchor_multilat():
    print("... Starting multilateration.")
    dist_dict = {}
    for anchor in rssi_dict:
        dist_dict[anchor] = dist_from_rssi(median(rssi_dict[anchor]))

    # https://github.com/kamalshadi/Localization
    import localization as lx
    P=lx.Project(mode='2D',solver='LSE')

    print("... adding anchors")
    for anchor in anchor_positions:
        P.add_anchor(anchor, anchor_positions[anchor])

    t,label = P.add_target()

    print("... adding measurements")
    for dist in dist_dict:
        P.add_measure(dist, dist_dict[dist])

    print("... calculating...")
    P.solve()

    print("> Done! Multilat result:", t.loc)

def broadcast_ping():
    rssi_dict = {}
    # Broadcast ping
    send_data("ping", "\xFF\xFF")

if __name__ == "__main__":
    sense = SenseHat()
    print(">> Opening serial port...")
    ser = serial.Serial("/dev/ttyUSB1", 9600)
    xbee = XBee(ser, callback=receive_data)

    rssi_list = []
    rssi_dict = {}

    ## TODO: Fill anchor position dictionary!
    anchor_positions = {
        "add1": (0, 1),
        "add2": (0, 3),
        "add3": (0, 4)
    }

    print(">> Waiting for events...")
    print("Middle: clear_matrix, left: init_rssi_calc, right: three_anchor_bbox, down: broadcast_ping")
    print("Sequence: broadcast_ping -> three_anchor_bbox / three_anchor_multilat")
    while True:
        for event in sense.stick.get_events():
            if event.action == "pressed":
                if event.direction == "middle":
                    print("** Event: Pressed.")
                    clear_matrix()
                elif event.direction == "left":
                    print("** Event: Left")
                    init_rssi_calc()
                elif event.direction == "right":
                    print("** Event: Right.")
                    three_anchor_bbox()
                elif event.direction == "down":
                    print("** Event: Down")
                    broadcast_ping()
                elif event.direction == "up":
                    print("** Event: up")
                    three_anchor_multilat()


import re
import sys
from napalm import get_network_driver
from getpass import getpass

def pretty_print(d, indent=0):
    for key, value in d.items():
        print('\t' * indent + str(key))
        if isinstance(value, dict):
            pretty_print(value, indent+1)
        elif isinstance(value, list):
            for e in value:
                print('\t' * (indent+1) + str(e))
        else:
            print('\t' * (indent+1) + str(value))

def build_dict(cfg):
    """    Builds nested/deep dictionary from Cisco ios
    config using recursion function 'child()', which also
    changes "the most child" dictionaries to list if possible.
    For global Cisco commands make special keys based on
    first word in the command, e.g.: '# aaa #'
    """
    def group_global_childless(dct):
        for k in list(dct):
            if not dct[k]:
                dct.pop(k,None)
                w = k.split()
                if w[0] == 'no':
                    sec_name = f"# {w[1]} #"
                else:
                    sec_name = f"# {w[0]} #"
                if sec_name in dct.keys():
                    dct[sec_name].append(k)
                else:
                    dct.update({sec_name: [k]})

    def child(base_indent):
        nonlocal n
        result = {}
        while True:
            if n >= len(lines):
                break
            stripped = lines[n].lstrip()
            indent = len(lines[n]) - len(stripped)
            if base_indent >= indent:
                break
            n = n + 1
            result.update({stripped: child(indent)})
        # In case we got all values={} transform result to list
        if not [v for v in result.values() if v]:
            result = [k for k in result.keys()]
        return result

    n = 0
    cfg, special_cases = cut_special_cases(cfg)
    lines = cfg.splitlines()
    lines = [line for line in lines if line
                and not line.startswith('!')
                and not line.startswith('end')]
    dct = child(base_indent=-1)
    dct.update(special_cases)
    group_global_childless(dct)
    return(dct)

def cut_special_cases(cfg):
    """ Cut special cases (banners, boot markers) from config and
    put them in special_cases dictionary, that is also returned
    """
    special_cases = {}
    rgx = r"((?:(?P<type>(?:set\s+)*banner\s\w+\s+)(?P<delim>\S+))((.*\r?\n)+?.*?)(\3).*)"
    re_banners = re.findall(rgx,cfg)
    for r in re_banners:
        cfg = cfg.replace(r[0],"",1)
        special_cases.update({f"# {r[1]}#": r[0].splitlines()})
    rgx = r"boot-start-marker\r?\n(.*\r?\n)*boot-end-marker"
    re_boot = re.search(rgx,cfg)
    cfg = cfg.replace(re_boot[0],"",1)
    special_cases.update({"# boot #": re_boot[0].splitlines()})
    return cfg, special_cases

def main():
    """    Reads config from file passed as argument. Or connect
    to Cisco ios device asking interactively ip, user, password
    Then prints result with indentation
    """
    if len(sys.argv) >= 1:
        file_name = sys.argv[1]
        fp = open(file_name)
        content = fp.read()
        dct = build_dict(content)
    else:
        driver = get_network_driver('ios')
        ipaddress = input('IP address: ')
        username = input('Username: ')
        ios_conn = driver(ipaddress, username, getpass())
        ios_conn.open()
        cfgs = ios_conn.get_config()
        dct = build_dict(cfgs['running'])
    pretty_print(dct,1)

if __name__ == "__main__":
    main()


from distutils.core import setup
from os.path import isdir
from itertools import product

all_packages = ['airborne_car_simulator']
packages = list(filter(isdir, all_packages))

setup(
    name='ese615',
    packages=packages,
    version='0.1',
    install_requires=[
            'pyyaml',
            'opencv-python',
            'matplotlib == 3.7.1',
            'numpy',
            'scipy'])