# -*- coding: UTF-8 -*-

from portality.lib import dates
from portality.datasets import get_country_code, get_currency_code
from copy import deepcopy
import locale, json, warnings
from urllib.parse import urlparse
from datetime import date, datetime

from portality.lib.dates import FMT_DATETIME_MS_STD, FMT_DATE_STD

#########################################################
## Data coerce functions

def to_currency_code(val):
    if val is None:
        return None
    nv = get_currency_code(val)
    if nv is None:
        raise ValueError("Unable to convert {x} to a valid currency code".format(x=val))
    uc = to_unicode()
    return uc(nv)

def to_country_code(val):
    if val is None:
        return None
    nv = get_country_code(val, fail_if_not_found=True)
    if nv is None:
        raise ValueError("Unable to convert {x} to a valid country code".format(x=val))
    uc = to_unicode()
    return uc(nv)

def to_unicode():
    def to_utf8_unicode(val):
        if isinstance(val, str):
            return val
        elif isinstance(val, str):
            try:
                return val.decode("utf8", "strict")
            except UnicodeDecodeError:
                raise ValueError("Could not decode string")
        else:
            return str(val)

    return to_utf8_unicode

def to_unicode_upper(val):
    ufn = to_unicode()
    val = ufn(val)
    return val.upper()

def to_int():
    def intify(val):
        # strip any characters that are outside the ascii range - they won't make up the int anyway
        # and this will get rid of things like strange currency marks
        if isinstance(val, str):
            val = val.encode("ascii", errors="ignore")

        # try the straight cast
        try:
            return int(val)
        except ValueError:
            pass

        # could have commas in it, so try stripping them
        try:
            return int(val.replace(",", ""))
        except ValueError:
            pass

        # try the locale-specific approach
        try:
            return locale.atoi(val)
        except ValueError:
            pass

        raise ValueError("Could not convert string to int: {x}".format(x=val))

    return intify

def to_float():
    def floatify(val):
        # strip any characters that are outside the ascii range - they won't make up the float anyway
        # and this will get rid of things like strange currency marks
        if isinstance(val, str):
            val = val.encode("ascii", errors="ignore")

        # try the straight cast
        try:
            return float(val)
        except ValueError:
            pass

        # could have commas in it, so try stripping them
        try:
            return float(val.replace(",", ""))
        except ValueError:
            pass

        # try the locale-specific approach
        try:
            return locale.atof(val)
        except ValueError:
            pass

        raise ValueError("Could not convert string to float: {x}".format(x=val))

    return floatify

def date_str(in_format=None, out_format=None):
    def datify(val):
        if val is None or val == "":
            return None
        if isinstance(val, date) or isinstance(val, datetime):
            return dates.format(val, format=out_format)
        else:
            return dates.reformat(val, in_format=in_format, out_format=out_format)

    return datify

def to_datestamp(in_format=None):
    def stampify(val):
        return dates.parse(val, format=in_format)

    return stampify

def to_isolang(output_format=None):
    """
    :param output_format: format from input source to putput.  Must be one of:
        * alpha3
        * alt3
        * alpha2
        * name
        * fr
    Can be a list in order of preference, too
    fixme: we could make these pycountry's keys, removing the need for so many transformations and intermediate steps
    :return:
    """
    # delayed import, since we may not always want to load the whole dataset for a dataobj
    from portality.lib import isolang as dataset

    # sort out the output format list
    if output_format is None:
        output_format = ["alpha3"]
    if not isinstance(output_format, list):
        output_format = [output_format]

    def isolang(val):
        if val is None:
            return None
        l = dataset.find(val)
        if l is None:
            raise ValueError("Unable to find iso code for language {x}".format(x=val))
        for f in output_format:
            v = l.get(f)
            if v is None or v == "":
                continue
            return v

    return isolang

def to_url(val):
    if not isinstance(val, str):
        raise ValueError("Argument passed to to_url was not a string, but type '{t}': '{val}'".format(t=type(val),val=val))

    val = val.strip()

    if val == '':
        return val

    # parse with urlparse
    url = urlparse(val)

    # now check the url has the minimum properties that we require
    if url.scheme and url.scheme.startswith("http"):
        uc = to_unicode()
        return uc(val)
    else:
        raise ValueError("Could not convert string {val} to viable URL".format(val=val))

def to_bool(val):
    """Conservative boolean cast - don't cast lists and objects to True, just existing booleans and strings."""
    if val is None:
        return None
    if val is True or val is False:
        return val

    if isinstance(val, str):
        if val.lower() == 'true':
            return True
        elif val.lower() == 'false':
            return False
        raise ValueError("Could not convert string {val} to boolean. Expecting string to either say 'true' or 'false' (not case-sensitive).".format(val=val))

    raise ValueError("Could not convert {val} to boolean. Expect either boolean or string.".format(val=val))

def string_canonicalise(canon, allow_fail=False):
    normalised = {}
    for a in canon:
        normalised[a.strip().lower()] = a

    def sn(val):
        if val is None:
            if allow_fail:
                return None
            raise ValueError("NoneType not permitted")

        try:
            norm = val.strip().lower()
        except:
            raise ValueError("Unable to treat value as a string")

        uc = to_unicode()
        if norm in normalised:
            return uc(normalised[norm])
        if allow_fail:
            return uc(val)

        raise ValueError("Unable to canonicalise string")

    return sn

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

############################################################
# The core data object which manages all the interactions
# with the underlying data member variable

class DataObjException(Exception):
    def __init__(self, *args, **kwargs):
        try:
            self.message = args[0]
        except IndexError:
            self.message = ''
        super(DataObjException, self).__init__(*args, **kwargs)

class DataSchemaException(DataObjException):
    pass

class DataObj(object):
    """
    Class which provides services to other classes which store their internal data
    as a python data structure in the self.data field.
    """

    SCHEMA = None

    DEFAULT_COERCE = {
        # NOTE - if you add something to the default coerce, add it to the default swagger
        # translation dict below as well. Furthermore if you're adding
        # custom stuff to the coerce, you will likely need to add an entry
        # to the swagger translation table as well, in the same way you
        # extend the coerce map.
        "unicode": to_unicode(),
        "unicode_upper" : to_unicode_upper,
        "utcdatetime": date_str(),
        "utcdatetimemicros" : date_str(out_format=FMT_DATETIME_MS_STD),
        "bigenddate" : date_str(out_format=FMT_DATE_STD),
        "integer": to_int(),
        "float": to_float(),
        "isolang": to_isolang(),
        "url": to_url,
        "bool": to_bool,
        "isolang_2letter": to_isolang(output_format="alpha2"),
        "country_code": to_country_code,
        "currency_code": to_currency_code,
        "license": string_canonicalise(["CC BY", "CC BY-NC", "CC BY-NC-ND", "CC BY-NC-SA", "CC BY-ND", "CC BY-SA", "Not CC-like"], allow_fail=True),
        "persistent_identifier_scheme": string_canonicalise(["None", "DOI", "Handles", "ARK"], allow_fail=True),
        "format": string_canonicalise(["PDF", "HTML", "ePUB", "XML"], allow_fail=True),
        "deposit_policy": string_canonicalise(["None", "Sherpa/Romeo", "Dulcinea", "OAKlist", "Diadorim"], allow_fail=True),
    }

    def __init__(self, raw=None, struct=None, construct_raw=True, expose_data=False, properties=None, coerce_map=None, construct_silent_prune=False, construct_maintain_reference=False, *args, **kwargs):
        # make a shortcut to the object.__getattribute__ function
        og = object.__getattribute__

        # if no subclass has set the coerce, then set it from default
        try:
            og(self, "_coerce_map")
        except:
            self._coerce_map = coerce_map if coerce_map is not None else deepcopy(self.DEFAULT_COERCE)

        # if no subclass has set the struct, initialise it
        try:
            og(self, "_struct")
        except:
            self._struct = struct

        # assign the data if not already assigned by subclass
        # NOTE: data is not _data deliberately
        try:
            og(self, "data")
        except:
            self.data = {} if raw is None else raw

        # properties to allow automatic object API construction
        # of the form
        #
        # {"<public property name>" : ("<path.to.property>", "<data object wrapper>")
        # e.g
        # {"identifier" : ("bibjson.identifier", DataObj))}
        try:
            og(self, "_properties")
        except:
            self._properties = properties if properties is not None else {}

        # if no subclass has set expose_data, set it
        try:
            og(self, "_expose_data")
        except:
            self._expose_data = expose_data

        # if no subclass has set _construct_silent_prune, set it
        try:
            og(self, "_construct_silent_prune")
        except:
            self._construct_silent_prune = construct_silent_prune

        # if no subclass has set _construct_maintain_reference, set it
        try:
            og(self, "_construct_maintain_reference")
        except:
            self._construct_maintain_reference = construct_maintain_reference

        # restructure the object based on the struct if requried
        if self._struct is not None and raw is not None and construct_raw:
            self.data = construct(self.data, self._struct, self._coerce_map, silent_prune=construct_silent_prune, maintain_reference=construct_maintain_reference)

        # run against the old validation routine
        # (now deprecated)
        self.validate()

        # run the object's native validation routine
        self.custom_validate()

        # keep a reference to the current data record, in case something up the inheritance chain messes with it
        # (I'm looking at you, UserDict).
        remember_this = self.data

        # finally, kick the request up
        super(DataObj, self).__init__(*args, **kwargs)
        self.data = remember_this

    def __getattr__(self, name):

        # workaround to prevent debugger from disconnecting at the deepcopy method
        # https://stackoverflow.com/questions/32831050/pycharms-debugger-gives-up-when-hitting-copy-deepcopy
        # if name.startwith("__"):
        #    raise AttributeError

        if hasattr(self.__class__, name):
            return object.__getattribute__(self, name)

        props, data_attrs = self._list_dynamic_properties()

        # if the name is not in the dynamic properties, raise an attribute error
        if name not in props and name not in data_attrs:
            raise AttributeError('{name} is not set'.format(name=name))

        # otherwise, extract the path from the properties list or the internal data
        if name in props:
            path, wrapper = self._properties.get(name)
        else:
            path = name
            wrapper = DataObj

        # request the internal property directly (which will in-turn raise the AttributeError if necessary)
        try:
            return self._get_internal_property(path, wrapper)
        except AttributeError:
            # re-wrap the attribute error with the name, rather than the path
            raise AttributeError('{name} is not set'.format(name=name))

    def __setattr__(self, key, value):
        # first set the attribute on any explicitly defined property
        try:
            if hasattr(self.__class__, key):
                # att = object.__getattribute__(self, key)
                return object.__setattr__(self, key, value)
        except AttributeError:
            pass

        # this could be an internal attribute from the constructor, so we need to make
        # a special case
        if key in ["_coerce_map", "_struct", "data", "_properties", "_expose_data"]:
            return object.__setattr__(self, key, value)

        props, data_attrs = self._list_dynamic_properties()

        # extract the path from the properties list or the internal data
        path = None
        wrapper = None
        if key in props:
            path, wrapper = self._properties.get(key)
        elif key in data_attrs:
            path = key
            wrapper = DataObj

        # try to set the property on othe internal object
        if path is not None:
            wasset = self._set_internal_property(path, value, wrapper)
            if wasset:
                return

        # fall back to the default approach of allowing any attribute to be set on the object
        return object.__setattr__(self, key, value)

    def check_construct(self):
        """
        Apply the construct to the internal data and throw errors if it is not validated

        This could be used, for example, if external processes have violated the .data encapsulation, or
        if internal processes which change .data need to be checked to make sure they haven't strayed outside
        their remit

        :return:
        """
        if self._struct is not None and self.data is not None:
            construct(self.data, self._struct, self._coerce_map, silent_prune=False, maintain_reference=False)

    def validate(self):
        """
        DEPRECATED - use 'check_construct' instead.

        :return:
        """
        warnings.warn("DEPRECATED - use 'check_construct' instead.", DeprecationWarning)
        if self.SCHEMA is not None:
            validate(self.data, self.SCHEMA)
        return True

    def custom_validate(self):
        pass

    def populate(self, fields_and_values):
        for k, v in fields_and_values.items():
            setattr(self, k, v)

    def clone(self):
        return self.__class__(deepcopy(self.data))

    def json(self):
        return json.dumps(self.data)

    def get_struct(self):
        return self._struct

    def _get_internal_property(self, path, wrapper=None):
        # pull the object from the structure, to find out what kind of retrieve it needs
        # (if there is a struct)
        type, substruct, instructions = None, None, None
        if self._struct:
            type, substruct, instructions = construct_lookup(path, self._struct)

        if type is None:
            # if there is no struct, or no object mapping was found, try to pull the path
            # as a single node (may be a field, list or dict, we'll find out in a mo)
            val = self._get_single(path)

            # if this is a dict or a list and a wrapper is supplied, wrap it
            if wrapper is not None:
                if isinstance(val, dict):
                    return wrapper(val, expose_data=self._expose_data)
                elif isinstance(val, list) and len(val) > 0:
                    if isinstance(val[0], dict):    # just check the first one
                        return [wrapper(v, expose_data=self._expose_data) for v in val]

            # otherwise, return the raw value if it is not None, or raise an AttributeError
            if val is None:
                raise AttributeError('{name} is not set'.format(name=path))

            return val

        if instructions is None:
            instructions = {}

        # if the struct contains a reference to the path, always return something, even if it is None - don't raise an AttributeError
        kwargs = construct_kwargs(type, "get", instructions)
        coerce_fn = self._coerce_map.get(instructions.get("coerce"))
        if coerce_fn is not None:
            kwargs["coerce"] = coerce_fn

        if type == "field":
            return self._get_single(path, **kwargs)
        elif type == "object":
            d = self._get_single(path, **kwargs)
            if wrapper:
                return wrapper(d, substruct, construct_raw=False, expose_data=self._expose_data)    # FIXME: this means all substructures are forced to use this classes expose_data policy, whatever it is
            else:
                return d
        elif type == "list":
            if instructions.get("contains") == "field":
                return self._get_list(path, **kwargs)
            elif instructions.get("contains") == "object":
                l = self._get_list(path, **kwargs)
                if wrapper:
                    return [wrapper(o, substruct, construct_raw=False, expose_data=self._expose_data) for o in l]    # FIXME: this means all substructures are forced to use this classes expose_data policy, whatever it is
                else:
                    return l

        # if for whatever reason we get here, raise the AttributeError
        raise AttributeError('{name} is not set'.format(name=path))

    def _set_internal_property(self, path, value, wrapper=None):

        def _wrap_validate(val, wrap, substruct):
            if wrap is None:
                if isinstance(val, DataObj):
                    return val.data
                else:
                    return val

            else:
                if isinstance(val, DataObj):
                    if isinstance(val, wrap):
                        return val.data
                    else:
                        raise AttributeError("Attempt to set {x} failed; is not of an allowed type.".format(x=path))
                else:
                    try:
                        d = wrap(val, substruct)
                        return d.data
                    except DataStructureException as e:
                        raise AttributeError(str(e))

        # pull the object from the structure, to find out what kind of retrieve it needs
        # (if there is a struct)
        type, substruct, instructions = None, None, None
        if self._struct:
            type, substruct, instructions = construct_lookup(path, self._struct)

        # if no type is found, then this means that either the struct was undefined, or the
        # path did not point to a valid point in the struct.  In the case that the struct was
        # defined, this means the property is trying to set something outside the struct, which
        # isn't allowed.  So, only set types which are None against objects which don't define
        # the struct.
        if type is None:
            if self._struct is None:
                if isinstance(value, list):
                    value = [_wrap_validate(v, wrapper, None) for v in value]
                    self._set_list(path, value)
                else:
                    value = _wrap_validate(value, wrapper, None)
                    self._set_single(path, value)

                return True
            else:
                return False

        if instructions is None:
            instructions = {}

        kwargs = construct_kwargs(type, "set", instructions)
        coerce_fn = self._coerce_map.get(instructions.get("coerce"))
        if coerce_fn is not None:
            kwargs["coerce"] = coerce_fn

        if type == "field":
            self._set_single(path, value, **kwargs)
            return True
        elif type == "object":
            v = _wrap_validate(value, wrapper, substruct)
            self._set_single(path, v, **kwargs)
            return True
        elif type == "list":
            if instructions.get("contains") == "field":
                self._set_list(path, value, **kwargs)
                return True
            elif instructions.get("contains") == "object":
                if not isinstance(value, list):
                    value = [value]
                vals = [_wrap_validate(v, wrapper, substruct) for v in value]
                self._set_list(path, vals, **kwargs)
                return True

        return False

    def _list_dynamic_properties(self):
        # list the dynamic properties the object could have
        props = []
        try:
            # props = og(self, 'properties').keys()
            props = list(self._properties.keys())
        except AttributeError:
            pass

        data_attrs = []
        try:
            if self._expose_data:
                if self._struct:
                    data_attrs = construct_data_keys(self._struct)
                else:
                    data_attrs = list(self.data.keys())
        except AttributeError:
            pass

        return props, data_attrs

    def _add_struct(self, struct):
        # if the struct is not yet set, set it
        try:
            object.__getattribute__(self, "_struct")
            self._struct = construct_merge(self._struct, struct)
        except:
            self._struct = struct

    def _get_path(self, path, default):
        parts = path.split(".")
        context = self.data

        for i in range(len(parts)):
            p = parts[i]
            d = {} if i < len(parts) - 1 else default
            context = context.get(p, d)
        return context

    def _set_path(self, path, val):
        parts = path.split(".")
        context = self.data

        for i in range(len(parts)):
            p = parts[i]

            if p not in context and i < len(parts) - 1:
                context[p] = {}
                context = context[p]
            elif p in context and i < len(parts) - 1:
                context = context[p]
            else:
                context[p] = val

    def _delete_from_list(self, path, val=None, matchsub=None, prune=True, apply_struct_on_matchsub=True):
        """
        Note that matchsub will be coerced with the struct if it exists, to ensure
        that the match is done correctly

        :param path:
        :param val:
        :param matchsub:
        :param prune:
        :return:
        """
        l = self._get_list(path)

        removes = []
        i = 0
        for entry in l:
            if val is not None:
                if entry == val:
                    removes.append(i)
            elif matchsub is not None:
                # attempt to coerce the sub
                if apply_struct_on_matchsub:
                    try:
                        object.__getattribute__(self, "_struct")
                        type, struct, instructions = construct_lookup(path, self._struct)
                        if struct is not None:
                            matchsub = construct(matchsub, struct, self._coerce_map)
                    except:
                        pass

                matches = 0
                for k, v in matchsub.items():
                    if entry.get(k) == v:
                        matches += 1
                if matches == len(list(matchsub.keys())):
                    removes.append(i)
            i += 1

        removes.sort(reverse=True)
        for r in removes:
            del l[r]

        if len(l) == 0 and prune:
            self._delete(path, prune)

    def _delete(self, path, prune=True):
        parts = path.split(".")
        context = self.data

        stack = []
        for i in range(len(parts)):
            p = parts[i]
            if p in context:
                if i < len(parts) - 1:
                    stack.append(context[p])
                    context = context[p]
                else:
                    del context[p]
                    if prune and len(stack) > 0:
                        stack.pop() # the last element was just deleted
                        self._prune_stack(stack)

    def _prune_stack(self, stack):
        while len(stack) > 0:
            context = stack.pop()
            todelete = []
            for k, v in context.items():
                if isinstance(v, dict) and len(list(v.keys())) == 0:
                    todelete.append(k)
            for d in todelete:
                del context[d]

    def _coerce(self, val, cast, accept_failure=False):
        if cast is None:
            return val
        try:
            return cast(val)
        except (ValueError, TypeError):
            if accept_failure:
                return val
            raise DataSchemaException("Cast with {x} failed on '{y}' of type {z}".format(x=cast, y=val, z=type(val)))

    def _get_single(self, path, coerce=None, default=None, allow_coerce_failure=True):
        # get the value at the point in the object
        val = self._get_path(path, default)

        if coerce is not None and val is not None:
            # if you want to coerce and there is something to coerce do it
            return self._coerce(val, coerce, accept_failure=allow_coerce_failure)
        else:
            # otherwise return the value
            return val

    def _get_list(self, path, coerce=None, by_reference=True, allow_coerce_failure=True):
        # get the value at the point in the object
        val = self._get_path(path, None)

        # if there is no value and we want to do by reference, then create it, bind it and return it
        if val is None and by_reference:
            mylist = []
            self._set_single(path, mylist)
            return mylist

        # otherwise, default is an empty list
        elif val is None and not by_reference:
            return []

        # check that the val is actually a list
        if not isinstance(val, list):
            raise DataSchemaException("Expecting a list at {x} but found {y}".format(x=path, y=val))

        # if there is a value, do we want to coerce each of them
        if coerce is not None:
            coerced = [self._coerce(v, coerce, accept_failure=allow_coerce_failure) for v in val]
            if by_reference:
                self._set_single(path, coerced)
            return coerced
        else:
            if by_reference:
                return val
            else:
                return deepcopy(val)

    def _set_single(self, path, val, coerce=None, allow_coerce_failure=False, allowed_values=None, allowed_range=None,
                    allow_none=True, ignore_none=False):

        if val is None and ignore_none:
            return

        if val is None and not allow_none:
            raise DataSchemaException("NoneType is not allowed at {x}".format(x=path))

        # first see if we need to coerce the value (and don't coerce None)
        if coerce is not None and val is not None:
            val = self._coerce(val, coerce, accept_failure=allow_coerce_failure)

        if allowed_values is not None and val not in allowed_values:
            raise DataSchemaException("Value {x} is not permitted at {y}".format(x=val, y=path))

        if allowed_range is not None:
            lower, upper = allowed_range
            if (lower is not None and val < lower) or (upper is not None and val > upper):
                raise DataSchemaException("Value {x} is outside the allowed range: {l} - {u}".format(x=val, l=lower, u=upper))

        # now set it at the path point in the object
        self._set_path(path, val)

    def _set_list(self, path, val, coerce=None, allow_coerce_failure=False, allow_none=True, ignore_none=False):
        # first ensure that the value is a list
        if not isinstance(val, list):
            val = [val]

        # now carry out the None check
        # for each supplied value, if it is none, and none is not allowed, raise an error if we do not
        # plan to ignore the nones.
        for v in val:
            if v is None and not allow_none:
                if not ignore_none:
                    raise DataSchemaException("NoneType is not allowed at {x}".format(x=path))

        # now coerce each of the values, stripping out Nones if necessary
        val = [self._coerce(v, coerce, accept_failure=allow_coerce_failure) for v in val if v is not None or not ignore_none]

        # check that the cleaned array isn't empty, and if it is behave appropriately
        if len(val) == 0:
            # this is equivalent to a None, so we need to decide what to do
            if ignore_none:
                # if we are ignoring nones, just do nothing
                return
            elif not allow_none:
                # if we are not ignoring nones, and not allowing them, raise an error
                raise DataSchemaException("Empty array not permitted at {x}".format(x=path))

        # now set it on the path
        self._set_path(path, val)

    def _add_to_list(self, path, val, coerce=None, allow_coerce_failure=False, allow_none=False, allowed_values=None, ignore_none=True, unique=False):
        if val is None and ignore_none:
            return

        if val is None and not allow_none:
            raise DataSchemaException("NoneType is not allowed in list at {x}".format(x=path))

        if allowed_values is not None and val not in allowed_values:
            raise DataSchemaException("Value {x} is not permitted at {y}".format(x=val, y=path))

        # first coerce the value
        if coerce is not None:
            val = self._coerce(val, coerce, accept_failure=allow_coerce_failure)
        current = self._get_list(path, by_reference=True)

        # if we require the list to be unique, check for the value first
        if unique:
            if val in current:
                return

        # otherwise, append
        current.append(val)

    def _set_with_struct(self, path, val):
        type, struct, instructions = construct_lookup(path, self._struct)
        if type == "field":
            kwargs = construct_kwargs(type, "set", instructions)
            coerce_fn = self._coerce_map.get(instructions.get("coerce", "unicode"))
            self._set_single(path, val, coerce=coerce_fn, **kwargs)
        elif type == "list":
            if not isinstance(val, list):
                val = [val]
            if struct is not None:
                val = [construct(x, struct, self._coerce_map) for x in val]
            kwargs = construct_kwargs(type, "set", instructions)
            coerce_fn = self._coerce_map.get(instructions.get("coerce"))
            self._set_list(path, val, coerce=coerce_fn, **kwargs)
        elif type == "object":
            if struct is not None:
                val = construct(val, struct, self._coerce_map)
            self._set_single(path, val)

    def _add_to_list_with_struct(self, path, val):
        type, struct, instructions = construct_lookup(path, self._struct)
        if type != "list":
            raise DataStructureException("Attempt to add to list {x} failed - it is not a list element".format(x=path))
        if struct is not None:
            val = construct(val, struct, self._coerce_map)
        kwargs = construct_kwargs(type, "set", instructions)
        self._add_to_list(path, val, **kwargs)

    def _utf8_unicode(self):
        """
        DEPRECATED - use dataobj.to_unicode() instead
        """
        return to_unicode()

    def _int(self):
        """
        DEPRECATED - use dataobj.to_int() instead
        """
        return to_int()

    def _float(self):
        """
        DEPRECATED - use dataobj.to_float() instead
        """
        return to_float()

    def _date_str(self, in_format=None, out_format=None):
        """
        DEPRECATED - use dataobj.date_str instead
        """
        return date_str(in_format=in_format, out_format=out_format)

############################################################
## Primitive object schema validation

class ObjectSchemaValidationError(DataObjException):
    pass

def validate(obj, schema):
    """
    DEPRECATED - use 'construct' instead.

    :param obj:
    :param schema:
    :return:
    """
    warnings.warn("DEPRECATED - use 'construct' instead.", DeprecationWarning)

    # all fields
    allowed = schema.get("bools", []) + schema.get("fields", []) + schema.get("lists", []) + schema.get("objects", [])

    for k, v in obj.items():
        # is k allowed at all
        if k not in allowed:
            raise ObjectSchemaValidationError("object contains key " + k + " which is not permitted by schema")

        # check the bools are bools
        if k in schema.get("bools", []):
            if type(v) != bool:
                raise ObjectSchemaValidationError("object contains " + k + " = " + str(v) + " but expected boolean")

        # check that the fields are plain old strings
        if k in schema.get("fields", []):
            if type(v) != str and type(v) != int and type(v) != float:
                raise ObjectSchemaValidationError("object contains " + k + " = " + str(v) + " but expected string, unicode or a number")

        # check that the lists are really lists
        if k in schema.get("lists", []):
            if type(v) != list:
                raise ObjectSchemaValidationError("object contains " + k + " = " + str(v) + " but expected list")
            # if it is a list, then for each member validate
            entry_schema = schema.get("list_entries", {}).get(k)
            if entry_schema is None:
                # validate the entries as fields
                for e in v:
                    if type(e) != str and type(e) != int and type(e) != float:
                        raise ObjectSchemaValidationError("list in object contains " + str(type(e)) + " but expected string, unicode or a number in " + k)
            else:
                # validate each entry against the schema
                for e in v:
                    validate(e, entry_schema)

        # check that the objects are objects
        if k in schema.get("objects", []):
            if type(v) != dict:
                raise ObjectSchemaValidationError("object contains " + k + " = " + str(v) + " but expected object/dict")
            # if it is an object, then validate
            object_schema = schema.get("object_entries", {}).get(k)
            if object_schema is None:
                #raise ObjectSchemaValidationError("no object entry for object " + k)
                pass # we are not imposing a schema on this object
            else:
                validate(v, object_schema)

############################################################
## Data structure coercion

class DataStructureException(DataObjException):
    pass

class ConstructException(DataObjException):
    pass

class ScriptTagFoundException(DataObjException):
    """
        Exception to raise when script tag has been found in metadata
    """
    pass

def construct_validate(struct, context=""):
    """
    Is the provided struct of the correct form
    {
        "fields" : {
            "field_name" : {"coerce" :"coerce_function", **kwargs}
        },
        "objects" : [
            "field_name"
        ],
        "lists" : {
            "field_name" : {"contains" : "object|field", "coerce" : "field_coerce_function, **kwargs}
        },
        "required" : ["field_name"],
        "structs" : {
            "field_name" : {
                <construct>
            }
        }
    }
    """
    # check that only the allowed keys are present
    keys = struct.keys()
    for k in keys:
        if k not in ["fields", "objects", "lists", "required", "structs"]:
            c = context if context != "" else "root"
            raise ConstructException("Key '{x}' present in struct at '{y}', but is not permitted".format(x=k, y=c))

    # now go through and make sure the fields are the right shape:
    for field_name, instructions in struct.get("fields", {}).items():
        if "coerce" not in instructions:
            c = context if context != "" else "root"
            raise ConstructException("Coerce function not listed in field '{x}' at '{y}'".format(x=field_name, y=c))
        for k,v in instructions.items():
            if not isinstance(v, list) and not isinstance(v, str):
                c = context if context != "" else "root"
                raise ConstructException("Argument '{a}' in field '{b}' at '{c}' is not a string or list".format(a=k, b=field_name, c=c))

    # then make sure the objects are ok
    for o in struct.get("objects", []):
        if not isinstance(o, str):
            c = context if context != "" else "root"
            raise ConstructException("There is a non-string value in the object list at '{y}'".format(y=c))

    # make sure the lists are correct
    for field_name, instructions in struct.get("lists", {}).items():
        contains = instructions.get("contains")
        if contains is None:
            c = context if context != "" else "root"
            raise ConstructException("No 'contains' argument in list definition for field '{x}' at '{y}'".format(x=field_name, y=c))
        if contains not in ["object", "field"]:
            c = context if context != "" else "root"
            raise ConstructException("'contains' argument in list '{x}' at '{y}' contains illegal value '{z}'".format(x=field_name, y=c, z=contains))
        for k,v in instructions.items():
            if not isinstance(v, list) and not isinstance(v, str):
                c = context if context != "" else "root"
                raise ConstructException("Argument '{a}' in list '{b}' at '{c}' is not a string or list".format(a=k, b=field_name, c=c))

    # make sure the requireds are correct
    for o in struct.get("required", []):
        if not isinstance(o, str):
            c = context if context != "" else "root"
            raise ConstructException("There is a non-string value in the required list at '{y}'".format(y=c))

    # now do the structs, which will involve some recursion
    substructs = struct.get("structs", {})

    # first check that there are no previously unknown keys in there
    possibles = struct.get("objects", []) + list(struct.get("lists", {}).keys())
    for s in substructs:
        if s not in possibles:
            c = context if context != "" else "root"
            raise ConstructException("struct contains key '{a}' which is not listed in object or list definitions at '{x}'".format(a=s, x=c))

    # now recurse into each struct
    for k,v in substructs.items():
        nc = context
        if nc == "":
            nc = k
        else:
            nc += "." + k
        construct_validate(v, context=nc)

    return True

def construct(obj, struct, coerce, context="", silent_prune=False, maintain_reference=False):
    """
    {
        "fields" : {
            "field_name" : {"coerce" :"coerce_function", **kwargs}

        },
        "objects" : [
            "field_name"
        ],
        "lists" : {
            "field_name" : {"contains" : "object|field", "coerce" : "field_coerce_function, **kwargs}
        },
        "required" : ["field_name"],
        "structs" : {
            "field_name" : {
                <construct>
            }
        }
    }

    :param obj:
    :param struct:
    :param coerce:
    :return:
    """
    if obj is None:
        return None

    # check that all the required fields are there
    try:
        keys = list(obj.keys())
    except:
        c = context if context != "" else "root"
        raise DataStructureException("Expected an object at {c} but found something else instead".format(c=c))

    for r in struct.get("required", []):
        if r not in keys:
            c = context if context != "" else "root"
            raise DataStructureException("Field '{r}' is required but not present at '{c}'".format(r=r, c=c))

    # check that there are no fields that are not allowed
    # Note that since the construct mechanism copies fields explicitly, silent_prune literally just turns off this
    # check
    if not silent_prune:
        allowed = list(struct.get("fields", {}).keys()) + struct.get("objects", []) + list(struct.get("lists", {}).keys())
        for k in keys:
            if k not in allowed:
                c = context if context != "" else "root"
                raise DataStructureException("Field '{k}' is not permitted at '{c}'".format(k=k, c=c))

    # this is the new object we'll be creating from the old
    constructed = DataObj()

    # now check all the fields
    for field_name, instructions in struct.get("fields", {}).items():
        val = obj.get(field_name)
        if val is None:
            continue
        coerce_fn = coerce.get(instructions.get("coerce", "unicode"))
        if coerce_fn is None:
            raise DataStructureException("No coersion function defined for type '{x}' at '{c}'".format(x=instructions.get("coerce", "unicode"), c=context + field_name))

        kwargs = construct_kwargs("field", "set", instructions)

        try:
            constructed._set_single(field_name, val, coerce=coerce_fn, **kwargs)
        except DataSchemaException as e:
            raise DataStructureException("Schema exception at '{a}', {b}".format(a=context + field_name, b=str(e)))

    # next check all the objetcs (which will involve a recursive call to this function)
    for field_name in struct.get("objects", []):
        val = obj.get(field_name)
        if val is None:
            continue
        if type(val) != dict:
            raise DataStructureException("Found '{x}' = '{y}' but expected object/dict".format(x=context + field_name, y=val))

        instructions = struct.get("structs", {}).get(field_name)

        if instructions is None:
            # this is the lowest point at which we have instructions, so just accept the data structure as-is
            # (taking a deep copy to destroy any references)
            try:
                constructed._set_single(field_name, deepcopy(val))
            except DataSchemaException as e:
                raise DataStructureException(str(e))
        else:
            # we need to recurse further down
            beneath = construct(val, instructions, coerce=coerce, context=context + field_name + ".", silent_prune=silent_prune)

            # what we get back is the correct sub-data structure, which we can then store
            try:
                constructed._set_single(field_name, beneath)
            except DataSchemaException as e:
                raise DataStructureException(str(e))

    # now check all the lists
    for field_name, instructions in struct.get("lists", {}).items():
        vals = obj.get(field_name)
        if vals is None:
            continue
        if not isinstance(vals, list):
            raise DataStructureException("Expecting list at {x} but found something else".format(x=context + field_name))

        # prep the keyword arguments for the setters
        kwargs = construct_kwargs("list", "set", instructions)

        contains = instructions.get("contains")
        if contains == "field":
            # coerce all the values in the list
            coerce_fn = coerce.get(instructions.get("coerce", "unicode"))
            if coerce_fn is None:
                raise DataStructureException("No coersion function defined for type '{x}' at '{c}'".format(x=instructions.get("coerce", "unicode"), c=context + field_name))

            for i in range(len(vals)):
                val = vals[i]
                try:
                    constructed._add_to_list(field_name, val, coerce=coerce_fn, **kwargs)
                except DataSchemaException as e:
                    raise DataStructureException(str(e))

        elif contains == "object":
            # for each object in the list, send it for construction
            for i in range(len(vals)):
                val = vals[i]

                if type(val) != dict:
                    raise DataStructureException("Found '{x}[{p}]' = '{y}' but expected object/dict".format(x=context + field_name, y=val, p=i))

                subinst = struct.get("structs", {}).get(field_name)
                if subinst is None:
                    try:
                        constructed._add_to_list(field_name, deepcopy(val))
                    except DataSchemaException as e:
                        raise DataStructureException(str(e))
                else:
                    # we need to recurse further down
                    beneath = construct(val, subinst, coerce=coerce, context=context + field_name + "[" + str(i) + "].", silent_prune=silent_prune)

                    # what we get back is the correct sub-data structure, which we can then store
                    try:
                        constructed._add_to_list(field_name, beneath)
                    except DataSchemaException as e:
                        raise DataStructureException(str(e))

        else:
            raise DataStructureException("Cannot understand structure where list '{x}' elements contain '{y}'".format(x=context + field_name, y=contains))

    if maintain_reference:
        obj.clear()
        obj.update(constructed.data)
        return obj
    else:
        return constructed.data

def construct_merge(target, source):
    merged = deepcopy(target)

    for field, instructions in source.get("fields", {}).items():
        if "fields" not in merged:
            merged["fields"] = {}
        if field not in merged["fields"]:
            merged["fields"][field] = deepcopy(instructions)

    for obj in source.get("objects", []):
        if "objects" not in merged:
            merged["objects"] = []
        if obj not in merged["objects"]:
            merged["objects"].append(obj)

    for field, instructions in source.get("lists", {}).items():
        if "lists" not in merged:
            merged["lists"] = {}
        if field not in merged["lists"]:
            merged["lists"][field] = deepcopy(instructions)

    for r in source.get("required", []):
        if "required" not in merged:
            merged["required"] = []
        if r not in merged["required"]:
            merged["required"].append(r)

    for field, struct in source.get("structs", {}).items():
        if "structs" not in merged:
            merged["structs"] = {}
        if field not in merged["structs"]:
            merged["structs"][field] = deepcopy(struct)
        else:
            # recursively merge
            merged["structs"][field] = construct_merge(merged["structs"][field], struct)

    return merged

def construct_lookup(path, struct):
    bits = path.split(".")

    # if there's more than one path element, we will need to recurse
    if len(bits) > 1:
        # it has to be an object, in order for the path to still have multiple
        # segments
        if bits[0] not in struct.get("objects", []):
            return None, None, None
        substruct = struct.get("structs", {}).get(bits[0])
        return construct_lookup(".".join(bits[1:]), substruct)
    elif len(bits) == 1:
        # first check the fields
        instructions = struct.get("fields", {}).get(bits[0])
        if instructions is not None:
            return "field", None, instructions

        # then check the lists
        instructions = struct.get("lists", {}).get(bits[0])
        if instructions is not None:
            structure = struct.get("structs", {}).get(bits[0])
            return "list", structure, instructions

        # then check the objects
        if bits[0] in struct.get("objects", []):
            structure = struct.get("structs", {}).get(bits[0])
            return "object", structure, None

    return None, None, None

def construct_kwargs(type, dir, instructions):
    # if there are no instructions there are no kwargs
    if instructions is None:
        return {}

    # take a copy of the instructions that we can modify
    kwargs = deepcopy(instructions)

    # remove the known arguments for the field type
    if type == "field":
        if "coerce" in kwargs:
            del kwargs["coerce"]

    elif type == "list":
        if "coerce" in kwargs:
            del kwargs["coerce"]
        if "contains" in kwargs:
            del kwargs["contains"]

    nk = {}
    if dir == "set":
        for k, v in kwargs.items():
            # basically everything is a "set" argument unless explicitly stated to be a "get" argument
            if not k.startswith("get__"):
                if k.startswith("set__"):    # if it starts with the set__ prefix, remove it
                    k = k[5:]
                nk[k] = v
    elif dir == "get":
        for k, v in kwargs.items():
            # must start with "get" argument
            if k.startswith("get__"):
                nk[k[5:]] = v

    return nk

def construct_data_keys(struct):
    return list(struct.get("fields", {})) + list(struct.get("objects", [])) + list(struct.get("lists", {}))

def merge_outside_construct(struct, target, source):
    merged = deepcopy(target)

    for source_key in source.keys():
        # if the source_key is one of the struct's fields, ignore it
        if source_key in list(struct.get("fields", {}).keys()):
            continue

        # if the source_key is one of the struct's lists, ignore it
        if source_key in list(struct.get("lists", {}).keys()):
            continue

        # if the source_key is one of the struct's object, we will need to go deeper
        if source_key in struct.get("objects", []):
            subsource = source[source_key]
            subtarget = target.get(source_key, {})
            substruct = struct.get("structs", {}).get(source_key, {})
            merged[source_key] = merge_outside_construct(substruct, subtarget, subsource)
            continue

        # if we get to here, the field in the source is not represented at this level in the struct,
        # so we should copy it over in full (unless the target already has a value here)
        if source_key not in merged:
            merged[source_key] = deepcopy(source[source_key])

    return merged

############################################################
## Unit test support

def test_dataobj(obj, fields_and_values):
    """
    Test a dataobj to make sure that the getters and setters you have specified
    are working correctly.

    Provide it a data object and a list of fields with the values to set and the expeceted return values (if required):

    {
        "key" : ("set value", "get value")
    }

    If you provide only the set value, then the get value will be required to be the same as the set value in the test

    {
        "key" : "set value"
    }

    :param obj:
    :param fields_and_values:
    :return:
    """
    for k, valtup in fields_and_values.items():
        if not isinstance(valtup, tuple):
            valtup = (valtup,)
        set_val = valtup[0]
        try:
            setattr(obj, k, set_val)
        except AttributeError:
            assert False, "Unable to set attribute {x} with value {y}".format(x=k, y=set_val)

    for k, valtup in fields_and_values.items():
        if not isinstance(valtup, tuple):
            valtup = (valtup,)
        get_val = valtup[0]
        if len(valtup) > 1:
            get_val = valtup[1]
        val = getattr(obj, k)
        assert val == get_val, (k, val, get_val)


#!/usr/local/bin/python
# -*- coding:utf-8 -*-
'''
 @Author:      py.gooker
 @Email:       zhangshch0131@126.com
 @DateTime:    2016-05-06 18:24:26
 @Description: Description
'''

import matplotlib.pyplot as plt
import sys

def save1():
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot([1, 2, 3])
    fig.savefig('test.png')

def save2():
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot([1, 2, 3])

    fig.savefig(sys.stdout)

if __name__ == '__main__':
    save1()
    # save2()


import torch
import wandb
import os

def denorm_visualizer(normed_image:torch.Tensor, mean=torch.Tensor([0.485, 0.456, 0.406]), std=torch.Tensor([0.229, 0.224, 0.225])):
    return ((normed_image.detach().permute(1, 2, 0).cpu() * std + mean)).numpy()

class WandbLogger(object):
    """ usage warning:this class only support scalar(_Tensor)
    create before training
    update(head, step, k=v)
    flush() to force writing
    """
    def __init__(self, args):
        os.environ["WANDB_API_KEY"] = "d32c431927c09a7bf392135ee5b63acaa90f5bee"
        if args.debug:
            os.environ["WANDB_MODE"] = "dryrun"
        self.is_master = (args.rank == 0)
        if self.is_master:
            wandb.init(project='cvq_vae_debug')
            wandb.config.update(args)
        self.class_labels = {
            0: 'background',
            1: 'aeroplane',
            2: 'bicycle',
            3: 'bird',
            4: 'boat',
            5: 'bottle',
            6: 'bus',
            7: 'car',
            8: 'cat',
            9: 'chair',
            10: 'cow',
            11: 'diningtable',
            12: 'dog',
            13: 'horse',
            14: 'motorbike',
            15: 'person',
            16: 'pottedplant',
            17: 'sheep',
            18: 'sofa',
            19: 'train',
            20: 'tvmonitor',
            255: 'void'
        }

    def update(self, **kwargs):
        if not self.is_master:
            return
        for k, v in kwargs.items():
            if v is None:
                continue
            if isinstance(v, tuple):
                wandb.log({k: wandb.Image(denorm_visualizer(v[0]),
                                    masks={
                                        f"mask_{i}": {"mask_data":v[i], "class_labels": self.class_labels}
                                    for i in range(1, len(v))})
                            })
            if isinstance(v, torch.Tensor):
                v = v.item()
            assert isinstance(v, (float, int))
            wandb.log({k:v})


from __future__ import unicode_literals

from unittest import TestCase

import ddt
import mock
from hamcrest import assert_that, equal_to, instance_of, raises

from storops.exception import UnityLunNameInUseError, JobStateError, \
    UnityPoolNameInUseError
from storops.unity.enums import RaidTypeEnum, FastVPStatusEnum, \
    FastVPRelocationRateEnum, PoolDataRelocationTypeEnum, \
    RaidStripeWidthEnum, TierTypeEnum, PoolUnitTypeEnum, \
    FSSupportedProtocolEnum, TieringPolicyEnum, JobStateEnum, \
    StoragePoolTypeEnum, ESXFilesystemMajorVersionEnum, \
    ESXFilesystemBlockSizeEnum, NFSShareDefaultAccessEnum, \
    NFSShareSecurityEnum
from storops.unity.resource.disk import UnityDiskGroup, UnityDisk
from storops.unity.resource.lun import UnityLun
from storops.unity.resource.nas_server import UnityNasServer
from storops.unity.resource.pool import UnityPool, UnityPoolList, \
    RaidGroupParameter
from storops.unity.resource.snap_schedule import UnitySnapSchedule
from storops.unity.resource.sp import UnityStorageProcessor
from storops.unity.resource.system import UnitySystem
from storops_test.unity.rest_mock import t_rest, patch_rest

__author__ = 'Cedric Zhuang'

@ddt.ddt
class UnityPoolTest(TestCase):
    @patch_rest
    def test_properties(self):
        pool = UnityPool(_id='pool_1', cli=t_rest())
        self.verify_pool_1(pool)

    @staticmethod
    def verify_pool_1(pool):
        assert_that(pool.id, equal_to('pool_1'))
        assert_that(pool.raid_type, equal_to(RaidTypeEnum.MIXED))
        assert_that(pool.name, equal_to('perfpool1130'))
        assert_that(pool.description, equal_to('pp'))
        assert_that(pool.size_free, equal_to(9160359936000))
        assert_that(pool.size_total, equal_to(9251627991040))
        assert_that(pool.size_used, equal_to(91268055040))
        assert_that(pool.size_subscribed, equal_to(1392106274816))
        assert_that(pool.alert_threshold, equal_to(70))
        assert_that(pool.pool_space_harvest_high_threshold, equal_to(95.0))
        assert_that(pool.pool_space_harvest_low_threshold, equal_to(85.0))
        assert_that(pool.snap_space_harvest_high_threshold, equal_to(25.0))
        assert_that(pool.snap_space_harvest_low_threshold, equal_to(20.0))
        assert_that(pool.is_fast_cache_enabled, equal_to(False))
        assert_that(str(pool.creation_time),
                    equal_to('2016-02-29 07:34:23+00:00'))
        assert_that(pool.is_empty, equal_to(False))
        assert_that(pool.is_harvest_enabled, equal_to(True))
        assert_that(pool.is_snap_harvest_enabled, equal_to(False))
        assert_that(pool.metadata_size_subscribed, equal_to(59324235776))
        assert_that(pool.snap_size_subscribed, equal_to(873220538368))
        assert_that(pool.metadata_size_used, equal_to(36775657472))
        assert_that(pool.snap_size_used, equal_to(24452407296))
        assert_that(pool.is_all_flash, equal_to(False))
        assert_that(pool.pool_type, equal_to(StoragePoolTypeEnum.TRADITIONAL))
        tiers = pool.tiers

        assert_that(len(tiers), equal_to(3))
        for tier in tiers:
            assert_that(tier._cli, equal_to(pool._cli))

        assert_that(pool.pool_fast_vp._cli, equal_to(pool._cli))

    @patch_rest
    def test_pool_fast_vp_properties(self):
        pool = UnityPool(_id='pool_1', cli=t_rest())
        fast = pool.pool_fast_vp
        assert_that(fast.status, equal_to(FastVPStatusEnum.ACTIVE))
        assert_that(fast.relocation_rate,
                    equal_to(FastVPRelocationRateEnum.MEDIUM))
        assert_that(fast.type, equal_to(PoolDataRelocationTypeEnum.SCHEDULED))
        assert_that(fast.is_schedule_enabled, equal_to(True))
        assert_that(str(fast.relocation_duration_estimate),
                    equal_to('0:00:00'))
        assert_that(fast.size_moving_down, equal_to(0))
        assert_that(fast.size_moving_up, equal_to(0))
        assert_that(fast.size_moving_within, equal_to(0))
        assert_that(fast.percent_complete, equal_to(0))
        assert_that(fast.data_relocated, equal_to(0))
        assert_that(str(fast.last_start_time),
                    equal_to('2016-03-13 22:00:00+00:00'))
        assert_that(str(fast.last_end_time),
                    equal_to('2016-03-14 06:00:00+00:00'))

    @patch_rest
    def test_tier_properties(self):
        pool = UnityPool(_id='pool_1', cli=t_rest())
        tier = next(t for t in pool.tiers if t.name == 'Performance')
        assert_that(tier.raid_type, equal_to(RaidTypeEnum.RAID5))
        assert_that(tier.stripe_width, equal_to(RaidStripeWidthEnum._5))
        assert_that(tier.tier_type, equal_to(TierTypeEnum.PERFORMANCE))
        assert_that(tier.size_total, equal_to(1180847570944))
        assert_that(tier.size_used, equal_to(3489660928))
        assert_that(tier.size_free, equal_to(1177357910016))
        assert_that(tier.size_moving_down, equal_to(0))
        assert_that(tier.size_moving_up, equal_to(0))
        assert_that(tier.size_moving_within, equal_to(0))
        assert_that(tier.disk_count, equal_to(5))

    @patch_rest
    def test_get_all(self):
        pools = UnityPoolList(cli=t_rest())
        assert_that(len(pools), equal_to(2))

        pool = next(pool for pool in pools if pool.id == 'pool_1')
        self.verify_pool_1(pool)

    @patch_rest
    def test_disk_groups(self):
        cli = t_rest()
        pool0 = UnityPool.get(cli=cli, _id='pool_1')
        disk_groups = pool0.disk_groups
        assert_that(disk_groups, instance_of(dict))
        assert_that(len(disk_groups), equal_to(2))
        assert_that(disk_groups['dg_8'], instance_of(list))
        assert_that(disk_groups['dg_15'], instance_of(list))
        for key in disk_groups:
            for disk in disk_groups[key]:
                assert_that(disk, instance_of(UnityDisk))

    @patch_rest
    def test_get_nested_resource_properties(self):
        pools = UnityPoolList(cli=t_rest())
        pool = next(pool for pool in pools if pool.id == 'pool_1')
        tier = next(t for t in pool.tiers if t.name == 'Performance')
        unit = next(u for u in tier.pool_units if u.id == 'rg_2')
        assert_that(unit.type, equal_to(PoolUnitTypeEnum.RAID_GROUP))
        assert_that(unit.tier_type, equal_to(TierTypeEnum.PERFORMANCE))
        assert_that(unit.name, equal_to("RAID5, #2, pool:perfpool1130"))
        assert_that(unit.description, equal_to('123'))
        assert_that(unit.wwn, equal_to(
            '06:00:00:00:05:00:00:00:01:00:00:00:00:00:00:64'))
        assert_that(unit.size_total, equal_to(1181501882368))
        assert_that(unit.pool, instance_of(UnityPool))

    @patch_rest
    def test_get_nested_resource_filter_by_non_id(self):
        pools = UnityPoolList(cli=t_rest())
        pool = next(pool for pool in pools if pool.id == 'pool_1')
        tier = next(t for t in pool.tiers if t.name == 'Performance')
        unit = next(u for u in tier.pool_units if u.description == '123')
        assert_that(unit.id, equal_to('rg_2'))

    @patch_rest
    def test_create_pool(self):
        cli = t_rest()
        disk_group = UnityDiskGroup.get(cli=cli, _id='dg_15')
        raid_group_0 = RaidGroupParameter(
            disk_group=disk_group,
            disk_num=3, raid_type=RaidTypeEnum.RAID5,
            stripe_width=RaidStripeWidthEnum.BEST_FIT)
        raid_groups = [raid_group_0]
        pool = UnityPool.create(
            cli=cli, name='test_pool', description='Unity test pool.',
            raid_groups=raid_groups, alert_threshold=15,
            is_harvest_enabled=True, is_snap_harvest_enabled=True,
            pool_harvest_high_threshold=80, pool_harvest_low_threshold=40,
            snap_harvest_high_threshold=80, snap_harvest_low_threshold=40,
            is_fast_cache_enabled=True, is_fastvp_enabled=True,
            pool_type=StoragePoolTypeEnum.DYNAMIC)
        assert_that(pool.id, equal_to('pool_4'))
        assert_that(pool.pool_type, equal_to(StoragePoolTypeEnum.DYNAMIC))
        assert_that(pool.is_all_flash, equal_to(False))

    @patch_rest
    def test_create_pool_name_in_use(self):
        cli = t_rest()
        raid_group_0 = RaidGroupParameter(
            disk_group='dg_15',
            disk_num=3, raid_type=RaidTypeEnum.RAID5,
            stripe_width=RaidStripeWidthEnum.BEST_FIT)
        raid_groups = [raid_group_0]

        def _inner():
            UnityPool.create(
                cli=cli, name='duplicate_pool',
                description='Unity test pool.',
                raid_groups=raid_groups)

        assert_that(_inner, raises(UnityPoolNameInUseError))

    @patch_rest
    def test_extend_pool(self):
        cli = t_rest()
        raid_group_0 = RaidGroupParameter(
            disk_group='dg_8',
            disk_num=4,
            raid_type=RaidTypeEnum.RAID10,
            stripe_width=RaidStripeWidthEnum.BEST_FIT)
        raid_groups = [raid_group_0]
        pool0 = UnityPool.get(cli=cli, _id='pool_1')
        resp = pool0.modify(raid_groups=raid_groups)
        assert_that(resp.is_ok(), equal_to(True))

    @patch_rest
    def test_modify_pool(self):
        cli = t_rest()
        pool0 = UnityPool.get(cli=cli, _id='pool_30')
        resp = pool0.modify(name="new_name",
                            is_fastvp_enabled=True,
                            alert_threshold=80)
        assert_that(resp.is_ok(), equal_to(True))

    @patch_rest
    def test_delete_pool(self):
        cli = t_rest()
        pool = UnityPool.get(cli=cli, _id='pool_4')
        resp = pool.delete()
        assert_that(resp.is_ok(), equal_to(True))

    @patch_rest
    def test_create_filesystem_success(self):
        pool = UnityPool(_id='pool_1', cli=t_rest())
        fs = pool.create_filesystem(
            'nas_2', 'fs3', 3 * 1024 ** 3,
            proto=FSSupportedProtocolEnum.CIFS,
            tiering_policy=TieringPolicyEnum.AUTOTIER_HIGH)
        assert_that(fs.get_id(), equal_to('fs_12'))

    @patch_rest
    def test_create_lun(self):
        pool = UnityPool(_id='pool_1', cli=t_rest())
        lun = pool.create_lun("LunName", 100)
        assert_that(lun, instance_of(UnityLun))

    @patch_rest
    def test_create_lun_with_same_name(self):
        pool = UnityPool(_id='pool_1', cli=t_rest())

        def f():
            pool.create_lun("openstack_lun")

        assert_that(f, raises(UnityLunNameInUseError))

    @patch_rest
    def test_create_lun_on_spb(self):
        pool = UnityPool(_id='pool_1', cli=t_rest())
        sp = UnityStorageProcessor(_id='spb', cli=t_rest())
        lun = pool.create_lun("LunName", 100, sp=sp)
        assert_that(lun, instance_of(UnityLun))

    @patch_rest
    def test_create_lun_with_muitl_property(self):
        pool = UnityPool(_id='pool_1', cli=t_rest())
        lun = pool.create_lun("LunName", 100,
                              description="Hello World", is_thin=True,
                              is_repl_dst=True,
                              tiering_policy=TieringPolicyEnum.AUTOTIER_HIGH)
        assert_that(lun, instance_of(UnityLun))

    @patch_rest
    def test_create_lun_with_compression_enabled_v4_2(self):
        pool = UnityPool(_id='pool_1', cli=t_rest(version='4.2'))
        lun = pool.create_lun("LunName", 100,
                              is_compression=True)
        assert_that(lun, instance_of(UnityLun))

    @patch_rest
    def test_create_lun_with_compression_enabled(self):
        pool = UnityPool(_id='pool_1', cli=t_rest(version='4.3'))
        lun = pool.create_lun("LunName", 100,
                              is_compression=True)
        assert_that(lun, instance_of(UnityLun))

    @patch_rest
    def test_create_lun_with_dedup_enabled(self):
        pool = UnityPool(_id='pool_1', cli=t_rest(version='5.0'))
        lun = pool.create_lun("LunName", 100,
                              is_compression=True,
                              is_advanced_dedup_enabled=True)
        assert_that(lun, instance_of(UnityLun))

    @patch_rest
    def test_create_vmfs(self):
        pool = UnityPool(_id='pool_1', cli=t_rest())
        vmfs = pool.create_vmfs(vmfs_name="VMFS datastore", size_gb=100)
        assert_that(vmfs, instance_of(UnityLun))

    @patch_rest
    def test_create_vmfs_vmware_iscsi_paramters(self):
        pool = UnityPool(_id='pool_1', cli=t_rest())
        vmfs = pool.create_vmfs(
            vmfs_name="VMFS datastore 2", size_gb=100,
            major_version=ESXFilesystemMajorVersionEnum.VMFS_5,
            block_size=ESXFilesystemBlockSizeEnum._2MB)
        assert_that(vmfs, instance_of(UnityLun))

    @patch_rest
    def test_create_nfs_share_success(self):
        pool = UnityPool(_id='pool_5', cli=t_rest())
        nas_server = UnityNasServer.get(cli=t_rest(), _id='nas_6')
        job = pool.create_nfs_share(
            nas_server,
            name='513dd8b0-2c22-4da0-888e-494d320303b6',
            size=4294967296)
        assert_that(JobStateEnum.COMPLETED, equal_to(job.state))

    @patch_rest
    def test_create_nfs_share_success_all_params(self):
        pool = UnityPool(_id='pool_5', cli=t_rest())
        nas_server = UnityNasServer.get(cli=t_rest(), _id='nas_6')
        size = 3 * 1024 ** 3

        job = pool.create_nfs_share(
            nas_server,
            name='513dd8b0-2c22-4da0-888e-494d320303b7',
            size=size, is_thin=True,
            tiering_policy=TieringPolicyEnum.AUTOTIER_HIGH,
            default_access=NFSShareDefaultAccessEnum.READ_WRITE,
            min_security=NFSShareSecurityEnum.KERBEROS,
            no_access_hosts_string='Host_1',
            read_only_hosts_string='Host_2',
            read_write_hosts_string='Host_3',
            read_only_root_hosts_string='Host_5,Host_4',
            root_access_hosts_string='Host_6',
            anonymous_uid=10001,
            anonymous_gid=10002,
            export_option=20001)
        assert_that(JobStateEnum.COMPLETED, equal_to(job.state))

    @patch_rest
    def test_create_nfs_share_failed(self):
        def f():
            pool = UnityPool(_id='pool_1', cli=t_rest())
            nas_server = UnityNasServer.get(cli=t_rest(), _id='nas_1')
            pool.create_nfs_share(
                nas_server,
                name='job_share_failed',
                size=1)

        assert_that(f, raises(JobStateError, 'too small'))

    @patch_rest
    def test_create_lun_with_snap_schedule(self):
        pool = UnityPool(_id='pool_1', cli=t_rest())
        schedule = UnitySnapSchedule(_id='snapSch_1', cli=t_rest())
        lun = pool.create_lun(
            lun_name='lun-with-snap-schedule',
            snap_schedule=schedule)
        assert_that(lun.get_id(), equal_to('sv_16455'))

    @patch_rest
    @ddt.data(
        {'unity_version': '4.5.0', 'pool_id': 'pool_1', 'is_all_flash': True,
         'expected': True},
        {'unity_version': '4.5.1', 'pool_id': 'pool_2', 'is_all_flash': False,
         'expected': False},
        {'unity_version': '5.0.3', 'pool_id': 'pool_1', 'is_all_flash': True,
         'expected': True},
        {'unity_version': '5.0.3', 'pool_id': 'pool_2', 'is_all_flash': False,
         'expected': False}
    )
    @ddt.unpack
    def test_is_compression_supported(self, unity_version, pool_id,
                                      is_all_flash, expected):
        cli = t_rest(unity_version)
        pool = UnityPool(_id=pool_id, cli=cli)
        pool.is_all_flash = is_all_flash
        assert_that(pool.is_compression_supported(), equal_to(expected))

    @patch_rest
    @ddt.data(
        {'unity_version': '4.3.0', 'unity_model': 'Unity 500',
         'pool_id': 'pool_1', 'is_all_flash': True, 'expected': False},
        {'unity_version': '4.3.1', 'unity_model': 'Unity 650F',
         'pool_id': 'pool_1', 'is_all_flash': True, 'expected': False},
        {'unity_version': '4.5.0', 'unity_model': 'Unity 500',
         'pool_id': 'pool_1', 'is_all_flash': True, 'expected': False},
        {'unity_version': '4.5.0', 'unity_model': 'Unity 500',
         'pool_id': 'pool_2', 'is_all_flash': False, 'expected': False},
        {'unity_version': '4.5.1', 'unity_model': 'Unity 650F',
         'pool_id': 'pool_1', 'is_all_flash': True, 'expected': True},
        {'unity_version': '4.5.1', 'unity_model': 'Unity 650F',
         'pool_id': 'pool_2', 'is_all_flash': False, 'expected': False},
        {'unity_version': '4.5.2', 'unity_model': 'Unity 640F',
         'pool_id': 'pool_2', 'is_all_flash': True, 'expected': False},
        {'unity_version': '4.5.2', 'unity_model': 'Unity 660F',
         'pool_id': 'pool_2', 'is_all_flash': True, 'expected': True},
        {'unity_version': '5.0.0', 'unity_model': 'Unity 500',
         'pool_id': 'pool_1', 'is_all_flash': True, 'expected': False},
        {'unity_version': '5.0.0', 'unity_model': 'Unity 500',
         'pool_id': 'pool_2', 'is_all_flash': False, 'expected': False},
        {'unity_version': '5.0.1', 'unity_model': 'Unity 480',
         'pool_id': 'pool_1', 'is_all_flash': True, 'expected': True},
        {'unity_version': '5.0.1', 'unity_model': 'Unity 480',
         'pool_id': 'pool_2', 'is_all_flash': False, 'expected': False},
        {'unity_version': '5.0.1', 'unity_model': 'Unity 470',
         'pool_id': 'pool_1', 'is_all_flash': True, 'expected': False},
        {'unity_version': '5.0.1', 'unity_model': 'Unity 490',
         'pool_id': 'pool_1', 'is_all_flash': True, 'expected': True},
        {'unity_version': '5.0.3', 'unity_model': 'Unity 880F',
         'pool_id': 'pool_1', 'is_all_flash': True, 'expected': True},
        {'unity_version': '5.0.3', 'unity_model': 'Unity 880F',
         'pool_id': 'pool_2', 'is_all_flash': False, 'expected': False},
    )
    @ddt.unpack
    def test_is_advanced_dedup_supported(self, unity_version, unity_model,
                                         pool_id, is_all_flash, expected):
        with mock.patch.object(UnitySystem, 'model', create=True,
                               return_value=unity_model,
                               new_callable=mock.PropertyMock):
            cli = t_rest(unity_version)
            pool = UnityPool(_id=pool_id, cli=cli)
            pool.is_all_flash = is_all_flash
            assert_that(pool.is_advanced_dedup_supported(), equal_to(expected))


from nltk.corpus import wordnet
from nltk import word_tokenize

def Overlap(sense1, sense2):
	def1 = set(word_tokenize(sense1.definition()))
	def2 = set(word_tokenize(sense2.definition()))
	intersect = len(def1.intersection(def2))
	total = len(def1) + len(def2) - intersect
	return float(intersect) / total

def Overlap_Based_Similarity(word1, word2):
	sense1 = wordnet.synsets(word1)
	sense2 = wordnet.synsets(word2)
	overlap = 0.0
	max_overlap= 0.0
	for i in sense1:
		for j in sense2:
			overlap = Overlap(i, j)
			if overlap > max_overlap:
				max_overlap = overlap
	return max_overlap

def main():
	word1 = input("Enter first word to check its similarity\t")
	word2 = input("Enter second word\t\t\t\t")

	sim = Overlap_Based_Similarity(word1.lower(), word2.lower())
	print("Overlap Based Similarity of " + word1 + " and "+ word2 + " is " + str(sim))

if __name__ == "__main__":
	main()


from collections import deque
import sys
input = sys.stdin.readline
delta = ((0,1),(0,-1),(1,0),(-1,0))

def bfs(r, c):
    q = deque()
    q.append((r, c))
    visited[r][c] = 1
    while q:
        r, c = q.popleft()
        for dr, dc in delta:
            nr, nc = dr+r, dc+c
            if 0 <= nr < H and 0 <= nc < W and not visited[nr][nc] and lst[nr][nc] == '#':
                q.append((nr, nc))
                visited[nr][nc] = 1

T = int(input())
for _ in range(T):
    H, W = map(int, input().split())
    lst = [list(input().rstrip()) for _ in range(H)]
    visited = [[0]*W for _ in range(H)]
    cnt = 0
    for i in range(H):
        for j in range(W):
            if lst[i][j] == '#' and not visited[i][j]:
                bfs(i, j)
                cnt += 1
    print(cnt)


import os, sys
from morl.hypervolume import InnerHyperVolume
def compute_hypervolume(ep_objs_batch):
    n = len(ep_objs_batch[0])
    HV = InnerHyperVolume(np.zeros(n))
    return HV.compute(ep_objs_batch)

base_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(base_dir)
sys.path.append(os.path.join(base_dir, 'externals/baselines'))
sys.path.append(os.path.join(base_dir, 'externals/pytorch-a2c-ppo-acktr-gail'))

from a2c_ppo_acktr import algo
from a2c_ppo_acktr.model import Policy

import pickle
import torch
import gym
import mujoco_py
import numpy as np
import environments
import argparse
from tqdm import tqdm
from collections import defaultdict
from morl.sample import Sample
from morl.population_2d import Population

parser = argparse.ArgumentParser()
parser.add_argument('--env_name', type=str, default='MO-Hopper-v2')
parser.add_argument('--collect_type', type=str, default="amateur")
parser.add_argument('--preference_type', type=str, default="uniform")
parser.add_argument('--num_traj', type=int, default=10000)
parser.add_argument('--data_path', type=str, default="data_collected")
parser.add_argument('--p_bar', type=bool, default=False)
args = parser.parse_args()

env_names = ["MO-Ant-v2", "MO-HalfCheetah-v2", "MO-Hopper-v2", "MO-Humanoid-v2", "MO-Swimmer-v2", "MO-Walker2d-v2"]
deterministic_transition = [True, True, True, True, True, True]
deterministic_initial = [True, True, True, False, True, True]
env_names_and_infos = {
    "MO-Ant-v2": {
        "base_name": "Ant-v2",
        "max_task": 69,
        "num_processes": 1,
        "base_kwargs": {'layernorm' : False},
        "obj_num": 2,
    },
    "MO-HalfCheetah-v2": {
        "base_name": "HalfCheetah-v2",
        "max_task": 233,
        "num_processes": 1,
        "base_kwargs": {'layernorm' : False},
        "obj_num": 2,
    },
    "MO-Hopper-v2": {
        "base_name": "Hopper-v2",
        "max_task": 95,
        "num_processes": 1,
        "base_kwargs": {'layernorm' : False},
        "obj_num": 2,
    },
    "MO-Humanoid-v2": {
        "base_name": "Humanoid-v2",
        "max_task": 354,
        "num_processes": 1,
        "base_kwargs": {'layernorm' : False},
        "obj_num": 2,
    },
    "MO-Swimmer-v2": {
        "base_name": "Swimmer-v2",
        "max_task": 173,
        "num_processes": 1,
        "base_kwargs": {'layernorm' : False},
        "obj_num": 2,
    },
    "MO-Walker2d-v2": {
        "base_name": "Walker2d-v2",
        "max_task": 364,
        "num_processes": 1,
        "base_kwargs": {'layernorm' : False},
        "obj_num": 2,
    },
    "MO-Hopper-v3": {
        "base_name": "Hopper-v3",
        "max_task": 2444,
        "num_processes": 1,
        "base_kwargs": {'layernorm': False},
        "obj_num": 3,
    }
}

def collect_helper(args, all_datas):
    data_path = f"{args.data_path}/{args.env_name}"
    if not os.path.exists(data_path):
        os.makedirs(data_path)
    filename = f"{data_path}/{args.env_name}_{args.num_traj}_new{args.collect_type}_{args.preference_type}.pkl"
    with open(filename, "wb") as f:
        pickle.dump(all_datas, f)

def generate_one_alpha(pref_dim, alpha_type):
    assert alpha_type in ['wide', 'narrow']
    _low = 0
    _high = 10e6
    alpha = np.random.uniform(low=_low, high=_high, size=pref_dim) if alpha_type == 'wide' \
        else np.random.uniform(low=_high/3, high=_high*2/3, size=pref_dim)
    return alpha

def dirichlet_preference(pref_dim, pref_range, n_alpha=1, n_pref_each_alpha=1):
    # sanity check
    assert pref_range in ['wide', 'narrow']
    pref_dim = int(pref_dim)
    n_alpha = int(n_alpha)
    n_pref_each_alpha = int(n_pref_each_alpha)
    # find alpha as uniform prior
    alphas = np.array([generate_one_alpha(pref_dim=pref_dim, alpha_type=pref_range) for _ in range(n_alpha)])
    # find preference using alpha + dirichlet
    preferences = np.zeros(shape=(n_alpha*n_pref_each_alpha, pref_dim))
    for i, alpha in enumerate(alphas):
        start = n_pref_each_alpha*i
        end = n_pref_each_alpha*(i+1)
        tmp = np.random.dirichlet(alpha, size=n_pref_each_alpha)
        tmp = tmp / np.linalg.norm(tmp, ord=1, axis=1, keepdims=True)
        preferences[start:end,:] = tmp
    return preferences

default_dirichlet_wide = lambda pref_dim, n_pref: dirichlet_preference(pref_dim, "wide", n_pref, 1)
defualt_dirichlet_narrow = lambda pref_dim, n_pref: dirichlet_preference(pref_dim, "narrow", n_pref, 1)

def uniform_preference(pref_dim, n_pref):
    k = np.array([np.random.exponential(scale=1.0, size=pref_dim) for _ in range(n_pref)])
    return k / np.sum(k, axis=1)[:, np.newaxis]

class PrefDist:

    def __init__(self, preference_type):
        if preference_type == "uniform":
            self.pref_func = uniform_preference
        elif preference_type == "wide":
            self.pref_func = default_dirichlet_wide
        elif preference_type == "narrow":
            self.pref_func = defualt_dirichlet_narrow

    def __call__(self, pref_dim, n_pref):
        return self.pref_func(pref_dim, n_pref)

class RejectSampling:

    def __init__(self, preference_dist, min_each_obj, max_each_obj, n_obj):
        self.preference_dist = preference_dist
        self.min_each_obj = min_each_obj
        self.max_each_obj = max_each_obj
        self.n_obj = n_obj

    def validate_preference(self, preference):
        for i, p in enumerate(preference):
            if p < self.min_each_obj[i] or p > self.max_each_obj[i]:
                return False # not valid
        return True # valid

    def get_preferences(self, n_pref_wanted):
        preferences = np.zeros(shape=(n_pref_wanted, n_obj))
        cur = 0
        while cur < n_pref_wanted:
            one_pref = self.preference_dist(n_obj, 1)[0]
            if self.validate_preference(one_pref):
                preferences[cur, :] = one_pref
                cur += 1
        return preferences

class GetPolicyId:

    def __init__(self, objs_normalized):
        self.objs_normalized = objs_normalized

    def __call__(self, sampled_preferences):
        policy_ids = np.zeros(shape=sampled_preferences.shape[0])
        for i, pref in enumerate(sampled_preferences):
            priorities = np.sum(np.square(self.objs_normalized - pref), axis=1) # sum of squared difference each row
            policy_ids[i] = np.argmin(priorities)
        return policy_ids
        # assume preference is valid here after reject-sampling

def update_item(datas, name, values):
    datas[name].append(values)

def eval_collect(args, samples, n_obj):

    all_datas = [{} for _ in range(args.num_traj)]
    eval_env = gym.make(args.env_name)
    objs_normalized = np.array([samp.objs_normalized for samp in samples])
    min_each_obj = np.min(objs_normalized, axis=0)
    max_each_obj = np.max(objs_normalized, axis=0)
    obj_range = max_each_obj - min_each_obj
    preference_dist = PrefDist(args.preference_type)

    if args.preference_type == "narrow":
        min_each_obj += obj_range / 3
        max_each_obj -= obj_range / 3

    if args.preference_type == "uniform":
        reject_sampler = RejectSampling(preference_dist, min_each_obj, max_each_obj, n_obj)
        sampled_preferences = reject_sampler.get_preferences(args.num_traj)
    else:
        sampled_preferences = preference_dist(n_obj, args.num_traj) * obj_range + min_each_obj
        sampled_preferences = sampled_preferences / np.sum(sampled_preferences, axis=1, keepdims=True)

    # PGMORL has a group of single-obj policies
    # here we try to decide the best single-obj to use for rollout
    getPolicyId = GetPolicyId(objs_normalized)
    policy_ids = getPolicyId(sampled_preferences)
    with torch.no_grad():

        for i in tqdm(range(args.num_traj), disable=(not args.p_bar)):

            preference = sampled_preferences[i, :]
            datas = defaultdict(list)

            policy_id = int(policy_ids[i])
            ob_rms = samples[policy_id].env_params['ob_rms']
            policy = samples[policy_id].actor_critic
            '''
            --------------------- EXPLANATION ---------------------
            we use RANDOMIZED environment reset instead of fixed reset value as in PGMORL
            this helps to diversify environment trajectories, so we need to give a seed
            '''
            eval_env.seed(i)
            obs = eval_env.reset()
            done = False
            ep_raw_reward = np.zeros(n_obj)

            while not done:

                update_item(datas, "observations", np.array(obs))
                # reload normalizing value used when training behavioral policy
                ob_norm = np.clip((obs - ob_rms.mean) / np.sqrt(ob_rms.var + 1e-8), -10.0, 10.0)

                # below should match in PEDA Appendix C: https://openreview.net/pdf?id=Ki4ocDm364
                if args.collect_type == "amateur":
                    if args.env_name in ['MO-Ant-v2', 'MO-Hopper-v2', 'MO-Hopper-v3', 'MO-Walker2d-v2', 'MO-HalfCheetah-v2']:
                        if np.random.uniform(0, 1) < 0.35:
                            action = policy.act(torch.Tensor(ob_norm).double().unsqueeze(0), None, None, deterministic=True)[1]
                        else:
                            action_deterministic = policy.act(torch.Tensor(ob_norm).double().unsqueeze(0), None, None, deterministic=True)[1]
                            action = action_deterministic * np.random.uniform(low=0.35, high=1.65)
                    elif args.env_name in ['MO-Swimmer-v2']:
                        if np.random.uniform(0, 1) < 0.35:
                            action = np.array([eval_env.action_space.sample()])
                        else:
                            action_deterministic = policy.act(torch.Tensor(ob_norm).double().unsqueeze(0), None, None, deterministic=True)[1]
                            action = action_deterministic * np.random.uniform(low=0.35, high=1.65)
                else: # expert
                    action = policy.act(torch.Tensor(ob_norm).double().unsqueeze(0), None, None, deterministic=True)[1]

                if args.env_name in ["MO-Hopper-v2", "MO-Hopper-v3"]:
                    action = np.clip(action, [-2, -2, -4], [2, 2, 4])
                else:
                    action = np.clip(action, -1, 1)

                next_obs, _, done, info = eval_env.step(action)

                raw_reward = info['obj']
                # reward = np.dot(preference, raw_reward)

                update_item(datas, "actions", np.array(action[0]))
                # update_item(datas, "rewards", np.array(reward))
                update_item(datas, "next_observations", np.array(next_obs))
                update_item(datas, 'preference', np.array(preference))
                update_item(datas, "terminals", np.array(done))
                update_item(datas, "raw_rewards", np.array(raw_reward))

                obs = next_obs
                ep_raw_reward += raw_reward

            for k, v in datas.items():
                datas[k] = np.array(v)

            all_datas[i] = datas

    eval_env.close()
    return all_datas

def eval_sample_hv(args, samples, eval_per_sample):

    with torch.no_grad():
        eval_out = np.zeros(shape=(len(samples), args.n_obj))
        for i, sample in enumerate(tqdm(samples, disable=(not args.p_bar))):
            policy = sample.actor_critic
            ob_rms = sample.env_params['ob_rms']

            all_this_sample = []
            for eval in range(eval_per_sample):
                env = gym.make(args.env_name)
                env.seed(np.random.randint(0, 100000))
                unweighted_raw_return = np.zeros(args.n_obj)
                done = False
                obs = env.reset()
                while not done:
                    ob_norm = np.clip((obs - ob_rms.mean) / np.sqrt(ob_rms.var + 1e-8), -10.0, 10.0)
                    action = policy.act(torch.Tensor(ob_norm).double().unsqueeze(0), None, None, deterministic=True)[1]
                    next_obs, _, done, info = env.step(action)
                    unweighted_raw_return += info['obj']
                    obs = next_obs
                all_this_sample.append(unweighted_raw_return)

            this_sample_median_ret = np.median(all_this_sample, axis=0)
            eval_out[i, :] = this_sample_median_ret
    hv = compute_hypervolume(eval_out)
    print(f"Env = {args.env_name}; data = {args.collect_type}_{args.preference_type}; hv = {hv:.3e}")

if __name__ == "__main__":

    env_name = args.env_name
    env_info = env_names_and_infos[env_name]
    base_env = gym.make(env_name)
    folder = f"Precomputed_Results/{env_info['base_name']}/final"

    objs_path = f"{folder}/objs.txt"
    with open(objs_path, 'r') as f:
        lines = f.read().splitlines()
        objs_original = np.array([[float(v) for v in l.split(",")] for l in lines])
    objs_normalized = objs_original / np.sum(objs_original, axis=1, keepdims=True)

    n_obj = env_info['obj_num']
    args.n_obj = n_obj
    samples = []
    for task in tqdm(range(env_info['max_task'] + 1), disable=(not args.p_bar)):

        policy_path = f"{folder}/EP_policy_{task}.pt"
        actor_critic = Policy(action_space=base_env.action_space,
                              obs_shape=base_env.observation_space.shape,
                              base_kwargs=env_info['base_kwargs'],
                              obj_num=n_obj)
        actor_critic.eval()
        actor_critic.load_state_dict(torch.load(policy_path))
        actor_critic.to("cpu").double()

        ob_rms_path = f"{folder}/EP_env_params_{task}.pkl"
        with open(ob_rms_path, 'rb') as f:
            ob = pickle.load(f)

        # default as given in PGMORL, only as a place-holder, these params doesn't
        # matter for collection
        agent = algo.PPO(
            actor_critic,
            None,
            10,
            32,
            0.5,
            0,
            lr=3e-4,
            eps=1e-5,
            max_grad_norm=None
        )

        sample = Sample(ob, actor_critic, agent)
        sample.objs = objs_original[task]
        sample.objs_normalized = objs_normalized[task]
        samples.append(sample)

    # also place holder here, these args are only needed for training
    # not for simulation
    args.pbuffer_num = 1
    args.pbuffer_size = 1

    '''
    optional: check the quality of behavioral policy ckpt (as rolled out again)
    '''
    # eval_sample_hv(args, samples, 5)

    all_datas = eval_collect(args, samples, n_obj)
    collect_helper(args, all_datas)


# -*- coding: utf-8 -*-
"""
(c) 2012-2022 Martin Wendt; see https://github.com/mar10/pyftpsync
Licensed under the MIT license: https://www.opensource.org/licenses/mit-license.php
"""

import os

import yaml

from ftpsync.cli_common import common_parser, creds_parser, verbose_parser
from ftpsync.synchronizers import CONFIG_FILE_NAME
from ftpsync.tree_command import tree_handler
from ftpsync.util import write

MANDATORY_TASK_ARGS = {"command", "remote"}
KNOWN_RUN_COMMANDS = {"upload", "download", "sync", "tree"}

KNOWN_TASK_ARGS = {
    "case",
    "create_folder",
    "debug",
    "delete_unmatched",
    "delete",
    "dry_run",
    "exclude",
    "files",  # tree command
    "force",
    "ftp_active",
    "here",
    "local",
    "match",
    "migrate",
    "no_color",
    "no_dry_run",  # alias: execute
    "no_keyring",
    "no_netrc",
    "no_prompt",
    "no_verify_host_keys",
    "progress",
    "prompt",
    "report_problems",
    "resolve",
    "root",
    "sort",  # tree command
    "verbose",
}

#: Boolean task options that can be overridden by passing an argument to the
#: run command. For example
#:     $ pyftpsync run --progress
#: will override the YAML setting
#:     tasks.TASK.progress: true
CLI_OVERRIDABLE_BOOL_TASK_ARGS = {
    "create_folder",
    "dry_run",
    "files",  # tree command
    "force",
    "no_color",
    "no_dry_run",
    "no_keyring",
    "no_netrc",
    "no_prompt",
    "no_verify_host_keys",
    "progress",
    "report_problems",
    "sort",  # tree command
}
_diff = CLI_OVERRIDABLE_BOOL_TASK_ARGS.difference(KNOWN_TASK_ARGS)
assert not _diff, f"Must be in KNOWN_TASK_ARGS: {_diff}"

def _set_default_task_arg(cli_args, task, name, default):
    """Add a task option to the CLI arguments namespace if not present."""
    if hasattr(cli_args, name):
        # write(f"SKIP: Yaml entry sets args.{name}: already {getattr(cli_args, name)}")
        return
    value = task.get(name, default)
    write(f"Yaml entry sets args.{name} => {value}")
    setattr(cli_args, name, value)

def add_run_parser(subparsers):
    # --- Create the parser for the "run" command ------------------------------

    parser = subparsers.add_parser(
        "run",
        parents=[verbose_parser, common_parser, creds_parser],
        help="run pyftpsync with configuration from `pyftpsync.yaml` in current or parent folder",
        allow_abbrev=False,
    )

    parser.add_argument(
        "task",
        nargs="?",
        help="task to run (default: use `default_task` from `pyftpsync.yaml`)",
    )

    p_group = parser.add_mutually_exclusive_group()
    p_group.add_argument(
        "--here", action="store_true", help="use current folder as root"
    )
    p_group.add_argument(
        "--root",
        action="store_true",
        help="use location of nearest `pyftpsync.yaml` as root folder",
    )

    parser.add_argument(
        "--no-dry-run",
        "--execute",  # alias
        action="store_true",
        help="override `dry_run: true` setting in `pyftpsync.yaml`",
    )
    parser.add_argument(
        "--create-folder",
        action="store_true",
        help="override `create_folder: false` or missing setting in `pyftpsync.yaml`",
    )
    parser.add_argument(
        "--files",
        action="store_true",
        help="override `files: false` or missing setting in `pyftpsync.yaml`",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="override `force: false` or missing setting in `pyftpsync.yaml`",
    )
    parser.add_argument(
        "--sort",
        action="store_true",
        help="override `sort: false` or missing setting in `pyftpsync.yaml`",
    )
    parser.add_argument(
        "--report-problems",
        action="store_true",
        help="override `report_problems: false` or missing setting in `pyftpsync.yaml`",
    )
    # TODO: --delete ?
    # TODO: --delete-unmatched ?
    # TODO: --no_keyring ?
    # TODO: --no_netrc ?
    # TODO: --no_prompt ?
    # TODO: --resolve ?

    parser.set_defaults(command="run")

    return parser

def handle_run_command(parser, cli_args):
    """Implement `run` sub-command."""
    MAX_LEVELS = 15

    # --- Look for `pyftpsync.yaml` in current folder and parents ---

    cur_level = 0
    cur_folder = os.getcwd()
    config_path = None
    while cur_level < MAX_LEVELS:
        path = os.path.join(cur_folder, CONFIG_FILE_NAME)
        # print("Searching for {}...".format(path))
        if os.path.isfile(path):
            config_path = path
            break
        parent = os.path.dirname(cur_folder)
        if parent == cur_folder:
            break
        cur_folder = parent
        cur_level += 1

    if not config_path:
        parser.error(
            f"Could not locate `pyftpsync.yaml` in {os.getcwd()} or {cur_level} parent folders."
        )

    # --- Parse and validate `pyftpsync.yaml` ---

    try:
        with open(config_path, "rt", encoding="utf-8-sig") as f:
            config = yaml.safe_load(f)
    except Exception as e:
        parser.error(f"Error parsing {config_path}: {e}")

    if "tasks" not in config:
        parser.error(f"Missing entry `tasks` in {config_path}")

    # --- Figure out which task to run ---

    if cli_args.task:
        task_name = cli_args.task
        if task_name not in config["tasks"]:
            parser.error(
                f"Invalid TASK argument {task_name!r}: "
                f"choose from [{', '.join(config['tasks'])}] "
                f"or define `tasks.{task_name}` in {config_path}"
            )

    elif config.get("default_task"):
        task_name = config.get("default_task")
        if task_name not in config["tasks"]:
            parser.error(
                f"Invalid entry `default_task: {task_name}`: "
                f"`tasks.{task_name}` is not defined in {config_path}"
            )

    else:
        parser.error(
            "No `TASK` argument was passed and no default configured: "
            f"choose from [{', '.join(config['tasks'])}] "
            f"or define `default_task: TASK` in {config_path}"
        )

    if task_name not in config["tasks"]:
        parser.error(f"Missing entry `tasks.{task_name}` in {config_path}")

    write(f"Running task '{task_name}' from {config_path}")

    # --- Task options inherit from common_config ---

    common_config = config.get("common_config", {})
    task = common_config.copy()
    task.update(config["tasks"][task_name])

    # --- Validate task configuration ---

    task_args = set(task.keys())

    missing_args = MANDATORY_TASK_ARGS.difference(task_args)
    if missing_args:
        parser.error(
            f"Missing mandatory entries: tasks.{task_name}.{', '.join(missing_args)}"
        )

    command = task["command"]
    if command not in KNOWN_RUN_COMMANDS:
        parser.error(
            f"Invalid entry: tasks.{task_name}.command: {command}: "
            f"expected one of {', '.join(KNOWN_RUN_COMMANDS)}"
        )

    allowed_args = KNOWN_TASK_ARGS.union(MANDATORY_TASK_ARGS)
    invalid_args = task_args.difference(allowed_args)
    if invalid_args:
        parser.error(f"Invalid entries: tasks.{task_name}.{', '.join(invalid_args)}")

    if "no_dry_run" in task:
        parser.error(
            "no_dry_run should not be used as yaml option: use dry_run: false instead"
        )

    if task.get("dry_run"):
        if cli_args.no_dry_run:
            write("`--no-dry-run` (or `--execute`) was passed: resetting dry_run mode")
        else:
            write(
                "Dry-run mode is configured: pass --no-dry-run (or --execute) "
                "to enable write operations"
            )

    # --- Override yaml entries by command line args ---

    for name in allowed_args:
        task_val = task.get(name, None)
        cli_val = getattr(cli_args, name, None)

        # write(
        #     f"check if entry `tasks.{task_name}.{name}: {task_val}` "
        #     f"is overriden by CLI arg `--{name}={cli_val!r}`"
        # )
        if cli_val != task_val:
            override = False
            if name in CLI_OVERRIDABLE_BOOL_TASK_ARGS and cli_val is True:
                override = True
                if name == "no_dry_run" and task.get("dry_run"):
                    # Already logged above
                    task["dry_run"] = False
            elif name in {"here", "root"} and (cli_args.here or cli_args.root):
                override = True
            elif name == "verbose" and cli_val != 3:
                assert type(cli_val) is int
                override = True

            if override:
                if cli_args.verbose >= 4:  # and name != "no_dry_run":
                    write(
                        f"Override yaml entry `tasks.{task_name}.{name}: {task_val}` "
                        f"with CLI arg `--{name}={cli_val!r}`"
                    )
                task[name] = cli_val

    # --- Add all configured and overridden task options to args namespace ---

    # Needed, because the subsequent command that is executed by this `run`
    # command may expect it:

    cli_args.command = task.get("command")

    for name in KNOWN_TASK_ARGS:
        task_val = task.get(name, None)
        cli_val = getattr(cli_args, name, None)
        if task_val is not None and task_val != cli_val:
            # write(f"Set args.{name} = {cli_val!r} => {task_val!r}")
            setattr(cli_args, name, task_val)

    # --- Figure out local target path ---

    cur_folder = os.getcwd()
    root_folder = os.path.dirname(config_path)
    if task.get("local"):
        root_folder = os.path.join(root_folder, task["local"])
    path_ofs = os.path.relpath(os.getcwd(), root_folder)

    if cur_level == 0 or task.get("root"):
        path_ofs = ""
        cli_args.local = root_folder
        cli_args.remote = task.get("remote")
    elif task.get("here"):
        write(f"Using sub-branch {path_ofs} of {root_folder}")
        cli_args.local = cur_folder
        cli_args.remote = os.path.join(task.get("remote"), path_ofs)
    else:
        parser.error(
            "`pyftpsync.yaml` configuration was found in a parent directory. "
            "Please pass an additional argument to clarify:\n"
            f"  --root: synchronize whole project ({root_folder})\n"
            f"  --here: synchronize sub branch ({root_folder}/{path_ofs})"
        )

    # --- Set the command handler and prepare args for subsequent call
    # --- of the configured task.command

    if command in ("upload", "download", "sync"):
        # This is handled by pyftpsync.run() afterwards
        pass
    # elif command == "scan":
    #     # Fix command line args as expected by command handler
    #     args.command = scan_handler
    #     args.target = args.remote
    #     _set_default_task_arg(args, task, "files", False)
    #     _set_default_task_arg(args, task, "sort", False)
    elif command == "tree":
        # Fix command line args as expected by command handler
        cli_args.command = tree_handler
        cli_args.target = cli_args.remote
        _set_default_task_arg(cli_args, task, "files", False)
        _set_default_task_arg(cli_args, task, "sort", False)
    else:
        raise NotImplementedError(f"Run task command {command}")