from typing import List, Optional, Dict, Any, Type
import logging
import logging.handlers
from datetime import datetime, timedelta
from glob import glob
import sys
import random
import warnings
import shutil
import os

import bugzoo
import cement
import yaml
import kaskara
from bugzoo.core import FileLine

from ..core import Language
from ..candidate import all_single_edit_patches
from ..candidate import Candidate
from ..transformation import Transformation
from ..transformation import find_all as find_all_transformations
from ..transformation.classic import DeleteStatement, \
                                     ReplaceStatement, \
                                     PrependStatement
from ..exceptions import BadConfigurationException, LanguageNotSupported
from ..searcher import Searcher
from ..problem import Problem
from ..localization import Localization, \
                           ample, \
                           genprog, \
                           jaccard, \
                           ochiai, \
                           tarantula
from ..snippet import SnippetDatabase
from ..settings import Settings
from .. import localization

logger = logging.getLogger(__name__)  # type: logging.Logger
logger.setLevel(logging.DEBUG)

BANNER = 'DARJEELING'

class BaseController(cement.Controller):
    class Meta:
        label = 'base'
        description = 'Language-independent automated program repair'
        arguments = [
            (['--version'], {'action': 'version', 'version': BANNER}),
        ]

    def default(self):
        # type: () -> None
        self.app.args.print_help()

    @property
    def _default_log_filename(self) -> str:
        # find all log file numbers that have been used in this directory
        used_numbers = [int(s.rpartition('.')[-1])
                        for s in glob('darjeeling.log.*')]

        if not used_numbers:
            return os.path.join(os.getcwd(), 'darjeeling.log.0')

        num = max(used_numbers) + 1
        return os.path.join(os.getcwd(), 'darjeeling.log.{}'.format(num))

    @cement.ex(
        help='attempt to automatically repair a given program',
        arguments=[
            (['filename'],
             {'help': ('a Darjeeling configuration file describing the faulty '
                       'program and how it should be repaired.') }),
            (['--log-to-file'],
             {'help': 'path to store the log file.',
              'type': str}),
            (['--seed'],
             {'help': 'random number generator seed',
              'type': int}),
            (['--max-candidates'],
             {'dest': 'limit_candidates',
              'type': int,
              'help': ('the maximum number of candidate patches that may be '
                       'considered by the search.')}),
            (['--max-time-mins'],
             {'dest': 'limit_time_minutes',
              'type': int,
              'help': ('the maximum number of minutes that may be spent '
                       'searching for a patch.')}),
            (['--continue'],
             {'dest': 'terminate_early',
              'action': 'store_false',
              'help': ('continue to search for patches after an acceptable '
                       ' patch has been discovered.')}),
            (['--threads'],
             {'dest': 'threads',
              'type': int,
              'help': ('number of threads over which the repair workload '
                       'should be distributed')})
        ]
    )
    def repair(self) -> None:
        # setup logging to file
        log_to_filename = self.app.pargs.log_to_file  # type: Optional[str]
        if not log_to_filename:
            log_to_filename = self._default_log_filename
        logger.info("logging to file: %s", log_to_filename)

        log_formatter = logging.Formatter(
            '%(asctime)s:%(name)s:%(levelname)s: %(message)s',
            '%Y-%m-%d %H:%M:%S')
        log_to_file = \
            logging.handlers.WatchedFileHandler(log_to_filename, mode='w')
        log_to_file.setLevel(logging.DEBUG)
        log_to_file.setFormatter(log_formatter)
        logging.getLogger('darjeeling').addHandler(log_to_file)

        filename = self.app.pargs.filename  # type: str
        seed = self.app.pargs.seed  # type: Optional[int]
        terminate_early = self.app.pargs.terminate_early  # type: bool
        threads = self.app.pargs.threads  # type: Optional[int]
        limit_candidates = \
            self.app.pargs.limit_candidates  # type: Optional[int]
        limit_time_minutes = \
            self.app.pargs.limit_time_minutes  # type: Optional[int]

        dir_patches = 'patches'
        if os.path.exists(dir_patches):
            logger.warning("destroying existing patch directory")
            shutil.rmtree(dir_patches)

        with open(filename, 'r') as f:
            yml = yaml.safe_load(f)

        has_limits = 'resource-limits' in yml

        # should we continue to search for repairs?
        if not terminate_early:
            logger.info("search will continue after an acceptable patch has been discovered")
        else:
            logger.info("search will terminate when an acceptable patch has been discovered")

        # how many threads should we use?
        if threads is not None:
            logger.info("using threads override: %d threads", threads)
        elif 'threads' in yml:
            if not isinstance(yml['threads'], int):
                m = "'threads' property should be an int"
                raise BadConfigurationException(m)
            threads = yml['threads']
            logger.info("using threads specified by configuration: %d threads",
                        threads)
        else:
            threads = 1
            logger.info("using default number of threads: %d", threads)
        if threads < 1:
            m = "number of threads must be greater than or equal to 1."
            raise BadConfigurationException(m)

        # determine time limit
        if limit_time_minutes is not None:
            logger.info("using time limit override: %d minutes",
                        limit_time_minutes)
        elif has_limits and 'time-minutes' in yml['resource-limits']:
            if not isinstance(yml['resource-limits']['time-minutes'], int):
                m = "'time-minutes' property in 'resource-limits' section should be an int"  # noqa: pycodestyle
                raise BadConfigurationException(m)
            limit_time_minutes = yml['resource-limits']['time-minutes']
            logger.info("using time limit specified by configuration: %d minutes",  # noqa: pycodestyle
                        limit_time_minutes)
        else:
            logger.info("no time limit is being enforced")

        limit_time = None  # type: Optional[timedelta]
        if limit_time_minutes:
            if limit_time_minutes < 1:
                m = "time limit must be greater than or equal to 1 minute"
                raise BadConfigurationException(m)
            limit_time = timedelta(minutes=limit_time_minutes)

        # determine the limit on the number of candidate repairs
        if limit_candidates is not None:
            logger.info("using candidate limit override: %d candidates",
                        limit_candidates)
        elif has_limits and 'candidates' in yml['resource-limits']:
            if not isinstance(yml['resource-limits']['candidates'], int):
                m = "'candidates' property in 'resource-limits' section should be an int"
                raise BadConfigurationException(m)
            limit_candidates = yml['resource-limits']['candidates']
            logger.info("using candidate limit specified by configuration: %d candidates",  # noqa: pycodestyle
                        limit_candidates)
        else:
            logger.info("no limit on number of candidate evaluations")

        if not limit_time and not limit_candidates:
            m = "no resource limits were specified; resource use will be unbounded"  # noqa: pycodestyle
            logger.warn(m)

        # seed override
        if seed is not None:
            logger.info("using random number generator seed override: %d",
                        seed)

        # no seed override; seed provided in config
        elif 'seed' in yml:
            if not isinstance(yml['seed'], int):
                m = "'seed' property should be an int."
                raise BadConfigurationException(m)
            elif yml['seed'] < 0:
                m = "'seed' property should be greater than or equal to zero."
                raise BadConfigurationException(m)
            seed = yml['seed']
            logger.info("using random number generator seed provided by configuration: %d",  # noqa: pycodestyle
                        seed)

        # no seed override or provided in provided
        # use current date/time
        elif seed is None:
            random.seed(datetime.now())
            seed = random.randint(0, sys.maxsize)
            logger.info("using random number generator seed based on current date and time: %d",  # noqa: pycodestyle
                        seed)

        # seed the RNG
        random.seed(seed)

        # determine the language
        if 'language' not in yml:
            m = "'language' property is missing"
            raise BadConfigurationException(m)
        if not isinstance(yml['language'], str):
            m = "'language' property should be a string"
            raise BadConfigurationException(m)
        try:
            language = Language.find(yml['language'])
        except LanguageNotSupported:
            supported = ', '.join([l.value for l in Language])
            supported = "(supported languages: {})".format(supported)
            m = "unsupported language [{}]. {}"
            m = m.format(yml['language'], supported)
            raise BadConfigurationException(m)
        logger.info("using language: %s", language.value)

        # build the settings
        opts = yml.get('optimizations', {})
        settings = \
            Settings(use_scope_checking=opts.get('use-scope-checking', True),
                     use_syntax_scope_checking=opts.get('use-syntax-scope-checking', True),
                     ignore_dead_code=opts.get('ignore-dead-code', True),
                     ignore_equivalent_appends=opts.get('ignore-equivalent-prepends', True),
                     ignore_untyped_returns=opts.get('ignore-untyped-returns', True),
                     ignore_string_equivalent_snippets=opts.get('ignore-string-equivalent-snippets', True),
                     ignore_decls=opts.get('ignore-decls', True),
                     only_insert_executed_code=opts.get('only-insert-executed-code', True))
        logger.info("using repair settings: %s", settings)

        # fetch the transformation schemas
        if 'transformations' not in yml:
            m = "'transformations' section is missing"
            raise BadConfigurationException(m)
        if not isinstance(yml['transformations'], dict):
            m = "'transformations' section should be an object"
            raise BadConfigurationException(m)
        if not 'schemas' in yml['transformations']:
            m = "'schemas' property missing in 'transformations' section"
            raise BadConfigurationException(m)
        if not isinstance(yml['transformations']['schemas'], list):
            m = "'schemas' property should be a list"
            raise BadConfigurationException(m)

        def schema_from_dict(d: Dict[str, Any]) -> Type[Transformation]:
            if not isinstance(d, dict):
                m = "expected an object but was a {}".format(type(d).__name__)
                m = "illegal schema description: {}".format(m)
                raise BadConfigurationException(m)
            if not 'type' in d:
                m = "missing 'type' property in schema description"
                raise BadConfigurationException(m)

            name_schema = d['type']
            try:
                return Transformation.find_schema(name_schema)
            except KeyError:
                known_schemas = "(known schemas: {})".format(
                    ', '.join(Transformation.schemas()))
                m = "no schema with name [{}]. {}"
                m = m.format(name_schema, known_schemas)
                raise BadConfigurationException(m)

        schemas = \
            [schema_from_dict(d) for d in yml['transformations']['schemas']]

        # fetch the bugzoo snapshot
        if 'snapshot' not in yml:
            raise BadConfigurationException("'snapshot' property is missing")
        if not isinstance(yml['snapshot'], str):
            m = "'snapshot' property should be a string"
            raise BadConfigurationException(m)
        name_snapshot = yml['snapshot']

        # connect to BugZoo
        logger.info("connecting to BugZoo server")
        with bugzoo.server.ephemeral(timeout_connection=120) as client_bugzoo:
            logger.info("connected to BugZoo server")
            try:
                snapshot = client_bugzoo.bugs[name_snapshot]
            except bugzoo.exceptions.BugZooException:
                logger.error("failed to fetch BugZoo snapshot: %s",
                             name_snapshot)
                sys.exit(1)

            if not client_bugzoo.bugs.is_installed(snapshot):
                logger.error("BugZoo snapshot is not installed: %s",
                             name_snapshot)
                sys.exit(1)

            # compute coverage
            logger.info("computing coverage information...")
            coverage = client_bugzoo.bugs.coverage(snapshot)
            logger.info("computed coverage information")

            # compute localization
            logger.info("computing fault localization...")
            if 'localization' not in yml:
                m = "'localization' section is missing"
                raise BadConfigurationException(m)

            localization = \
                Localization.from_config(coverage, yml['localization'])
            logger.info("computed fault localization:\n%s", localization)

            # determine implicated files and lines
            files = localization.files
            lines = list(localization)  # type: List[FileLine]

            # compute analysis
            analysis = kaskara.Analysis.build(client_bugzoo,
                                              snapshot,
                                              files)

            # build problem
            problem = Problem(bz=client_bugzoo,
                              bug=snapshot,
                              language=language,
                              coverage=coverage,
                              analysis=analysis,
                              settings=settings)

            # build snippet database
            logger.info("constructing database of donor snippets...")
            snippets = SnippetDatabase.from_statements(
                analysis.statements,
                use_canonical_form=settings.ignore_string_equivalent_snippets)
            logger.info("constructed database of donor snippets: %d snippets",
                        len(snippets))

            # FIXME build and index transformations
            logger.info("constructing transformation database...")
            tx = list(find_all_transformations(problem, lines, snippets, schemas))
            logger.info("constructed transformation database: %d transformations",  # noqa: pycodestyle
                        len(tx))

            # build the search strategy
            # FIXME pass limits!
            searcher = Searcher.from_dict(yml['algorithm'], problem, tx,
                                          threads=threads,
                                          candidate_limit=limit_candidates,
                                          time_limit=limit_time)

            logger.info("beginning search process...")
            patches = []  # type: List[Candidate]
            if terminate_early:
                try:
                    patches.append(next(searcher.__iter__()))
                except StopIteration:
                    pass
            else:
                patches = list(searcher)
            if not patches:
                logger.info("failed to find a patch")

            # wait for threads to finish gracefully before exiting
            searcher.close()

            # report stats
            num_test_evals = searcher.num_test_evals
            num_candidate_evals = searcher.num_candidate_evals
            time_running_mins = searcher.time_running.seconds / 60

            logger.info("found %d plausible patches", len(patches))
            logger.info("time taken: %.2f minutes", time_running_mins)
            logger.info("# test evaluations: %d", searcher.num_test_evals)
            logger.info("# candidate evaluations: %d", searcher.num_candidate_evals)

            # save patches to disk
            os.makedirs(dir_patches, exist_ok=True)
            for i, patch in enumerate(patches):
                diff = str(patch.to_diff(problem))
                fn_patch = os.path.join(dir_patches, '{}.diff'.format(i))
                logger.debug("writing patch to %s", fn_patch)
                try:
                    with open(fn_patch, 'w') as f:
                        f.write(diff)
                except Exception:
                    logger.exception("failed to write patch: %s", fn_patch)
                    raise
                logger.debug("wrote patch to %s", fn_patch)

class CLI(cement.App):
    class Meta:
        label = 'darjeeling'
        handlers = [BaseController]

def main():
    log_to_stdout = logging.StreamHandler()
    log_to_stdout.setLevel(logging.INFO)
    # logger.addHandler(log_to_stdout)
    logging.getLogger('darjeeling').addHandler(log_to_stdout)

    with CLI() as app:
        app.run()


#!/usr/bin/env python

#adapted from Conference Central by Wesley Chun(wesc+api@google.com)

from datetime import datetime

import endpoints
from protorpc import messages
from protorpc import message_types
from protorpc import remote

from google.appengine.api import memcache
from google.appengine.api import taskqueue
from google.appengine.ext import ndb

from models import ConflictException
from models import Profile
from models import ProfileForm
from models import ProfileMiniForm
from models import StringMessage
from models import BooleanMessage
from models import Activity
from models import ActivityForm
from models import ActivityForms
from models import ActivityQueryForm
from models import Tribe
from models import TribeForm
from models import TribeForms

from settings import WEB_CLIENT_ID
from settings import ANDROID_CLIENT_ID
from settings import IOS_CLIENT_ID
from settings import ANDROID_AUDIENCE

from utils import getUserId

EMAIL_SCOPE = endpoints.EMAIL_SCOPE
API_EXPLORER_CLIENT_ID = endpoints.API_EXPLORER_CLIENT_ID

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

DEFAULTS = {
    "time": "Time not specified",
    "approvalNeeded": False,
    "currentParticipants": 0,
}

ACTI_GET_REQUEST = endpoints.ResourceContainer(
    message_types.VoidMessage,
    websafeActivityKey=messages.StringField(1),
)

ACTI_POST_REQUEST = endpoints.ResourceContainer(
    ActivityForm,
    websafeActivityKey=messages.StringField(1),
)

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

@endpoints.api(name='activity', version='v1', audiences=[ANDROID_AUDIENCE],
    allowed_client_ids=[WEB_CLIENT_ID, API_EXPLORER_CLIENT_ID, ANDROID_CLIENT_ID, IOS_CLIENT_ID],
    scopes=[EMAIL_SCOPE])
class ActivityApi(remote.Service):

# - - - Tribe objects - - - - - - - - - - - - - - - - - -

    def _createTribeObject(self, request):
        """Create or update Tribe object, returning TribeForm/request."""
        data = (getattr(request, "tribename")).lower()
        tribe = Tribe(key=ndb.Key(Tribe, data),tribename=data)
        tribe.put()
        return request

    def _copyTribeToForm(self, tribe):
        """Copy relevant fields from Activity to ActivityForm."""
        tf = TribeForm()
        for field in tf.all_fields():
            setattr(tf, field.name, (getattr(tribe, field.name)))
        return tf

    @endpoints.method(TribeForm, TribeForm, path='tribe',
            http_method='POST', name='createTribe')
    def createTribe(self, request):
        """Create new tribe."""
        return self._createTribeObject(request)

    @endpoints.method(message_types.VoidMessage, TribeForms,
            path='queryTribes',
            http_method='POST',
            name='queryTribes')
    def queryTribes(self, request):
        """Query for activities."""
        tribes = Tribe.query()

        # return individual TribeForm object per Tribe
        return TribeForms(
                items=[self._copyTribeToForm(tribe) for tribe in \
                tribes]
        )

# - - - Activity objects - - - - - - - - - - - - - - - - -

    def _copyActivityToForm(self, acti, displayName):
        """Copy relevant fields from Activity to ActivityForm."""
        af = ActivityForm()
        for field in af.all_fields():
            if hasattr(acti, field.name):
                # convert Date to date string; just copy others
                if field.name=='date':
                    setattr(af, field.name, str(getattr(acti, field.name)))
                elif field.name=='participants':
                    setattr(af, field.name, (",".join(getattr(acti, field.name))))
                else:
                    setattr(af, field.name, (getattr(acti, field.name)))
            elif field.name == "websafeKey":
                setattr(af, field.name, acti.key.urlsafe())
        if displayName:
            setattr(af, 'organizerDisplayName', displayName)
        af.check_initialized()
        return af

    def _createActivityObject(self, request):
        """Create or update Activity object, returning ActivityForm/request."""
        # preload necessary data items
        user = endpoints.get_current_user()
        if not user:
            raise endpoints.UnauthorizedException('Authorization required')
        user_id = getUserId(user)

        if not request.tribe:
            raise endpoints.BadRequestException("Activity 'tribe' field required")

        # copy ActivityForm/ProtoRPC Message into dict
        data = {field.name: getattr(request, field.name) for field in request.all_fields()}
        del data['websafeKey']
        del data['organizerDisplayName']

        # add default values for those missing (both data model & outbound Message)
        for df in DEFAULTS:
            if data[df] in (None, []):
                data[df] = DEFAULTS[df]
                setattr(request, df, DEFAULTS[df])
        # convert dates from strings to Date objects
        if data['date']:
            data['date'] = datetime.strptime(data['date'][:10], "%Y-%m-%d").date()
        if data['time']:
            data['time'] = str(data['time'])+" "+data['ampm']
        del data['ampm']
        # generate Profile Key based on user ID and Activity
        # ID based on Profile key get Activity key from ID
        p_key = ndb.Key(Profile, user_id)
        a_id = Activity.allocate_ids(size=1, parent=p_key)[0]
        a_key = ndb.Key(Activity, a_id, parent=p_key)
        data['key'] = a_key
        data['organizerUID'] = user_id
        data['participants'] = []
         # creation of Activity & return (modified) ActivityForm
        Activity(**data).put()
        return request

    @ndb.transactional()
    def _updateActivityObject(self, request):
        user = endpoints.get_current_user()
        if not user:
            raise endpoints.UnauthorizedException('Authorization required')
        user_id = getUserId(user)

        # copy ActivityForm/ProtoRPC Message into dict
        data = {field.name: getattr(request, field.name) for field in request.all_fields()}

        # update existing activity
        acti = ndb.Key(urlsafe=request.websafeActivityKey).get()
        # check that activity exists
        if not acti:
            raise endpoints.NotFoundException(
                'No activity found with key: %s' % request.websafeActivityKey)

        # check that user is owner
        if user_id != acti.organizerUID:
            raise endpoints.ForbiddenException(
                'Only the owner can update the activity.')

        for field in request.all_fields():
            data = getattr(request, field.name)
            # only copy fields where there is new data
            if data not in (None, []):
                  # write to Activity object
                setattr(acti, field.name, data)
        acti.put()
        prof = ndb.Key(Profile, user_id).get()
        return self._copyActivityToForm(acti, getattr(prof, 'displayName'))

    @ndb.transactional(xg=True)
    def _activityRegistration(self, request, reg=True):
        """Register or unregister user for selected activity."""
        retval = None
        prof = self._getProfileFromUser() # get user Profile

        # check if acti exists given websafeActiKey
        # get activity; check that it exists
        wsak = request.websafeActivityKey
        acti = ndb.Key(urlsafe=wsak).get()
        if not acti:
            raise endpoints.NotFoundException(
                'No activity found with key: %s' % wsak)

        # register
        if reg:
            # check if user already registered otherwise add
            if wsak in prof.activityKeysToAttend:
                raise ConflictException(
                    "You have already registered for this activity")

            # check if seats avail
            if acti.maxParticipants==acti.currentParticipants:
                raise ConflictException(
                    "There are no places available.")

            # register user, take away one seat
            prof.activityKeysToAttend.append(wsak)
            acti.currentParticipants += 1
            acti.participants.append(prof.displayName)
            retval = True

        # unregister
        else:
            # check if user already registered
            if wsak in prof.activityKeysToAttend:

                # unregister user, add back one seat
                prof.activityKeysToAttend.remove(wsak)
                acti.currentParticipants -= 1
                acti.participants.remove(prof.displayName)
                retval = True
            else:
                retval = False

        # write things back to the datastore & return
        prof.put()
        acti.put()
        return BooleanMessage(data=retval)

    @endpoints.method(ActivityForm, ActivityForm, path='activity',
            http_method='POST', name='createActivity')
    def createActivity(self, request):
        """Create new activity."""
        return self._createActivityObject(request)

    @endpoints.method(ACTI_POST_REQUEST, ActivityForm,
            path='activity/{websafeActivityKey}',
            http_method='PUT', name='updateActivity')
    def updateActivity(self, request):
        """Update activity w/provided fields & return w/updated info."""
        return self._updateActivityObject(request)

    @endpoints.method(ACTI_GET_REQUEST, ActivityForm,
            path='activity/{websafeActivityKey}',
            http_method='GET', name='getActivity')
    def getActivity(self, request):
        """Return requested activity (by websafeActivityKey)."""
        # get Activity object from request; bail if not found
        acti = ndb.Key(urlsafe=request.websafeActivityKey).get()
        if not acti:
            raise endpoints.NotFoundException(
                'No activity found with key: %s' % request.websafeActivityKey)
        prof = acti.key.parent().get()
        # return ActivityForm
        return self._copyActivityToForm(acti, getattr(prof, 'displayName'))

    @endpoints.method(message_types.VoidMessage, ActivityForms,
            path='getActivitiesCreated',
            http_method='POST', name='getActivitiesCreated')
    def getActivitiesCreated(self, request):
        """Return activities created by user."""
        # make sure user is authed
        user = endpoints.get_current_user()
        if not user:
            raise endpoints.UnauthorizedException('Authorization required')
        user_id = getUserId(user)

        # create ancestor query for all key matches for this user
        actis = Activity.query(ancestor=ndb.Key(Profile, user_id))
        prof = ndb.Key(Profile, user_id).get()
        # return set of ActivityForm objects per Activity
        return ActivityForms(
            items=[self._copyActivityToForm(acti, getattr(prof, 'displayName')) for acti in actis]
        )

    @endpoints.method(message_types.VoidMessage, ActivityForms,
            path='activities/attending',
            http_method='GET', name='getActivitiesToAttend')
    def getActivitiesToAttend(self, request):
        """Get list of activities that user has registered for."""
        prof = self._getProfileFromUser() # get user Profile
        acti_keys = [ndb.Key(urlsafe=wsak) for wsak in prof.activityKeysToAttend]
        activities = ndb.get_multi(acti_keys)

        # get organizers
        organisers = [ndb.Key(Profile, acti.organizerUID) for acti in activities]
        profiles = ndb.get_multi(organisers)

        # put display names in a dict for easier fetching
        names = {}
        for profile in profiles:
            names[profile.key.id()] = profile.displayName

        # return set of ActivityForm objects per Activity
        return ActivityForms(items=[self._copyActivityToForm(acti, names[acti.organizerUID])\
         for acti in activities]
        )

    @endpoints.method(ACTI_GET_REQUEST, BooleanMessage,
            path='activity/{websafeActivityKey}',
            http_method='POST', name='registerForActivity')
    def registerForActivity(self, request):
        """Register user for selected activity."""
        return self._activityRegistration(request)

    @endpoints.method(ACTI_GET_REQUEST, BooleanMessage,
            path='activity/{websafeActivityKey}',
            http_method='DELETE', name='unregisterFromActivity')
    def unregisterFromActivity(self, request):
        """Unregister user for selected activity."""
        return self._activityRegistration(request, reg=False)

    @endpoints.method(ActivityQueryForm, ActivityForms,
            path='queryActivities',
            http_method='POST',
            name='queryActivities')
    def queryActivities(self, request):
        """Query for activities."""
        activities = Activity.query()
        if request.fil:
            activities = Activity.query(Activity.tribe==request.fil)
        # need to fetch organiser displayName from profiles
        # get all keys and use get_multi for speed
        organisers = [(ndb.Key(Profile, acti.organizerUID)) for acti in activities]
        profiles = ndb.get_multi(organisers)

        # put display names in a dict for easier fetching
        names = {}
        for profile in profiles:
            names[profile.key.id()] = profile.displayName

        # return individual ActivityForm object per Activity
        return ActivityForms(
                items=[self._copyActivityToForm(acti, names[acti.organizerUID]) for acti in \
                activities]
        )

# - - - Profile objects - - - - - - - - - - - - - - - - - - -

    def _copyProfileToForm(self, prof):
        """Copy relevant fields from Profile to ProfileForm."""
        # copy relevant fields from Profile to ProfileForm
        pf = ProfileForm()
        for field in pf.all_fields():
            if hasattr(prof, field.name):
                setattr(pf, field.name, getattr(prof, field.name))
        pf.check_initialized()
        return pf

    def _getProfileFromUser(self):
        """Return user Profile from datastore, creating new one if non-existent."""
        # make sure user is authed
        user = endpoints.get_current_user()
        if not user:
            raise endpoints.UnauthorizedException('Authorization required')

        # get Profile from datastore
        user_id = getUserId(user)
        p_key = ndb.Key(Profile, user_id)
        profile = p_key.get()
        # create new Profile if not there
        if not profile:
            profile = Profile(
                key = p_key,
                displayName = user.nickname(),
                mainEmail= user.email(),
            )
            profile.put()
        # return Profile
        return profile

    def _doProfile(self, save_request=None):
        """Get user Profile and return to user, possibly updating it first."""
        # get user Profile
        prof = self._getProfileFromUser()

        # if saveProfile(), process user-modifyable fields
        if save_request:
            for field in ('displayName', 'phoneNumber'):
                if hasattr(save_request, field):
                    val = getattr(save_request, field)
                    if val:
                        setattr(prof, field, (val))
                        prof.put()

        # return ProfileForm
        return self._copyProfileToForm(prof)

    @endpoints.method(message_types.VoidMessage, ProfileForm,
            path='profile', http_method='GET', name='getProfile')
    def getProfile(self, request):
        """Return user profile."""
        return self._doProfile()

    @endpoints.method(ProfileMiniForm, ProfileForm,
            path='profile', http_method='POST', name='saveProfile')
    def saveProfile(self, request):
        """Update & return user profile."""
        return self._doProfile(request)

# - - - Tribes - - - - - - - - - - - - - - - - - - - -

api = endpoints.api_server([ActivityApi])


# Viết một chương trình có thể tính giai thừa của một số cho trước.
# Kết quả được in thành chuỗi trên một dòng, phân tách bởi dấu phẩy.
# Ví dụ, số cho trước là 8 thì kết quả đầu ra phải là 40320.

res = 1

n = int(input('Nhap so cho truoc: '))

# for i in range(1, n + 1):
#    res *= i

def gt(x):
    if x == 1:
        return x
    return x * gt(x - 1)

res = gt(n)

print(res)


import copy

import pandas as pd
from stix2 import Filter
from itertools import chain
from tqdm import tqdm
import datetime
import re
import numpy as np

# Lookup module for Platforms - each matrix has a list of possible platforms, and each platform with multiple
#   subplatforms has a corresponding entry. This allows for a pseudo-recursive lookup of subplatforms, as the presence
#   of a platform at the top level of this lookup indicates the existence of subplatforms.
MATRIX_PLATFORMS_LOOKUP = {"enterprise-attack": ['PRE', 'Windows', 'macOS', 'Linux', 'Cloud', 'Office 365', 'Azure AD',
                                                 'Google Workspace', 'SaaS', 'IaaS', 'Network', 'Containers'],
                           "mobile-attack": ['Android', 'iOS'],
                           "Cloud": ['Office 365', 'Azure AD', 'Google Workspace', 'SaaS', 'IaaS'],
                           "ics-attack": ["Field Controller/RTU/PLC/IED", "Safety Instrumented System/Protection Relay",
                                          "Control Server", "Input/Output Server", "Windows", "Human-Machine Interface",
                                          "Engineering Workstation", "Data Historian"]}

def remove_revoked_deprecated(stix_objects):
    """Remove any revoked or deprecated objects from queries made to the data source"""
    # Note we use .get() because the property may not be present in the JSON data. The default is False
    # if the property is not set.
    return list(
        filter(
            lambda x: x.get("x_mitre_deprecated", False) is False and x.get("revoked", False) is False,
            stix_objects
        )
    )

def filter_platforms(stix_objects, platforms):
    """Filter out any objects that don't have a matching platform to one in 'platforms'"""
    if not platforms:
        return stix_objects

    return list(
        filter(
            lambda x: any(platform.lower() in [y.lower() for y in x.get("x_mitre_platforms", [])]
                          for platform in platforms),
            stix_objects
        )
    )

def format_date(date):
    """ Given a date string, return it formatted as %d %B %Y """
    if isinstance(date, str):
        date = datetime.datetime.strptime(date, "%Y-%m-%dT%H:%M:%S.%fZ")
    return "{} {} {}".format(date.strftime("%d"), date.strftime("%B"), date.strftime("%Y"))

def get_citations(objects):
    """given a list of STIX objects, return a pandas dataframe for the citations on the objects"""
    citations = []
    for sdo in objects:
        if "external_references" in sdo:
            for ref in sdo["external_references"]:
                if "external_id" not in ref and "description" in ref and not ref["description"].startswith(
                        "(Citation: "):
                    citation = {
                        "reference": ref["source_name"],
                        "citation": ref["description"]
                    }
                    if "url" in ref:
                        citation["url"] = ref["url"]

                    citations.append(citation)

    return pd.DataFrame(citations).drop_duplicates(subset="reference", ignore_index=True)

def parseBaseStix(sdo):
    """given an SDO, return a dict of field names:values that are common across all ATT&CK STIX types"""
    row = {}
    url = None
    if "external_references" in sdo and sdo["external_references"][0]["source_name"] in ["mitre-attack",
                                                                                "mitre-mobile-attack",
                                                                                "mitre-ics-attack"]:
        row["ID"] = sdo["external_references"][0]["external_id"]
        url = sdo["external_references"][0]["url"]
    if "name" in sdo:
        row["name"] = sdo["name"]
    if "description" in sdo:
        row["description"] = sdo["description"]
    if url:
        row["url"] = url
    if "created" in sdo:
        row["created"] = format_date(sdo["created"])
    if "modified" in sdo:
        row["last modified"] = format_date(sdo["modified"])
    if "x_mitre_version" in sdo:
        row["version"] = sdo["x_mitre_version"]
    if "x_mitre_contributors" in sdo:
        row["contributors"] = "; ".join(sorted(sdo["x_mitre_contributors"]))
    return row

def techniquesToDf(src, domain):
    """
    Parse STIX techniques from the given data and return corresponding pandas dataframes.
    :param src: MemoryStore or other stix2 DataSource object holding the domain data
    :param domain: domain of ATT&CK src corresponds to, e.g "enterprise-attack"
    :returns: a lookup of labels (descriptors/names) to dataframes
    """

    techniques = src.query([Filter("type", "=", "attack-pattern")])
    techniques = remove_revoked_deprecated(techniques)
    technique_rows = []

    for technique in tqdm(techniques, desc="parsing techniques"):
        # get parent technique if sub-technique
        subtechnique = "x_mitre_is_subtechnique" in technique and technique["x_mitre_is_subtechnique"]
        if subtechnique:
            subtechnique_of = src.query([
                Filter("type", "=", "relationship"),
                Filter("relationship_type", "=", "subtechnique-of"),
                Filter("source_ref", "=", technique["id"])
            ])[0]
            parent = src.get(subtechnique_of["target_ref"])

        # base STIX properties
        row = parseBaseStix(technique)

        # sub-technique properties
        tactic_shortnames = list(map(lambda kcp: kcp["phase_name"], technique["kill_chain_phases"]))
        tactics = list(map(lambda t: t.replace("-", " ").title(), tactic_shortnames))
        row["tactics"] = ", ".join(sorted(tactics))

        if "x_mitre_detection" in technique:
            row["detection"] = technique["x_mitre_detection"]
        if "x_mitre_platforms" in technique:
            row["platforms"] = ", ".join(sorted(technique["x_mitre_platforms"]))

        # domain specific fields -- ICS + Enterprise
        if domain in ["enterprise-attack", "ics-attack"]:
            if "x_mitre_data_sources" in technique:
                row["data sources"] = ", ".join(sorted(technique["x_mitre_data_sources"]))

        # domain specific fields -- enterprise
        if domain == "enterprise-attack":
            row["is sub-technique"] = subtechnique
            if subtechnique:
                row["name"] = f"{parent['name']}: {technique['name']}"
                row["sub-technique of"] = parent["external_references"][0]["external_id"]

            if "x_mitre_system_requirements" in technique:
                row["system requirements"] = ", ".join(sorted(technique["x_mitre_system_requirements"]))
            if "privilege-escalation" in tactic_shortnames:
                if "x_mitre_permissions_required" in technique:
                    row["permissions required"] = ", ".join(sorted(technique["x_mitre_permissions_required"]))
                if "x_mitre_effective_permissions" in technique:
                    row["effective permissions"] = ", ".join(sorted(technique["x_mitre_effective_permissions"]))
            if "defense-evasion" in tactic_shortnames and "x_mitre_defense_bypassed" in technique:
                row["defenses bypassed"] = ", ".join(sorted(technique["x_mitre_defense_bypassed"]))
            if "execution" in tactic_shortnames and "x_mitre_remote_support" in technique:
                row["supports remote"] = technique["x_mitre_remote_support"]
            if "impact" in tactic_shortnames and "x_mitre_impact_type" in technique:
                row["impact type"] = ", ".join(sorted(technique["x_mitre_impact_type"]))

        # domain specific fields -- mobile
        elif domain == "mobile-attack":
            if "x_mitre_tactic_type" in technique:
                row["tactic type"] = ", ".join(sorted(technique["x_mitre_tactic_type"]))
            mtc_refs = list(filter(lambda ref: ref["source_name"] == "NIST Mobile Threat Catalogue",
                                   technique["external_references"]))
            if mtc_refs:
                row["MTC ID"] = mtc_refs[0]["external_id"]

        technique_rows.append(row)

    citations = get_citations(techniques)
    dataframes = {
        "techniques": pd.DataFrame(technique_rows).sort_values("name"),
    }
    # add relationships
    dataframes.update(relationshipsToDf(src, relatedType="technique"))
    # add/merge citations
    if not citations.empty:
        if "citations" in dataframes:  # append to existing citations from references
            dataframes["citations"] = dataframes["citations"].append(citations)
        else:  # add citations
            dataframes["citations"] = citations

        dataframes["citations"].sort_values("reference")

    return dataframes

def tacticsToDf(src, domain):
    """
    Parse STIX tactics from the given data and return corresponding pandas dataframes.
    :param src: MemoryStore or other stix2 DataSource object holding the domain data
    :param domain: domain of ATT&CK src corresponds to, e.g "enterprise-attack"
    :returns: a lookup of labels (descriptors/names) to dataframes
    """

    tactics = src.query([Filter("type", "=", "x-mitre-tactic")])
    tactics = remove_revoked_deprecated(tactics)

    tactic_rows = []
    for tactic in tqdm(tactics, desc="parsing mitigations"):
        tactic_rows.append(parseBaseStix(tactic))

    citations = get_citations(tactics)
    dataframes = {
        "tactics": pd.DataFrame(tactic_rows).sort_values("name"),
    }
    if not citations.empty: dataframes["citations"] = citations.sort_values("reference")

    return dataframes

def softwareToDf(src, domain):
    """
    Parse STIX software from the given data and return corresponding pandas dataframes.
    :param src: MemoryStore or other stix2 DataSource object holding the domain data
    :param domain: domain of ATT&CK src corresponds to, e.g "enterprise-attack"
    :returns: a lookup of labels (descriptors/names) to dataframes
    """

    software = list(chain.from_iterable(  # software are the union of the tool and malware types
        src.query(f) for f in [Filter("type", "=", "tool"), Filter("type", "=", "malware")]
    ))
    software = remove_revoked_deprecated(software)
    software_rows = []
    for soft in tqdm(software, desc="parsing software"):
        # add common STIx fields
        row = parseBaseStix(soft)
        # add software-specific fields
        if "x_mitre_platforms" in soft:
            row["platforms"] = ", ".join(sorted(soft["x_mitre_platforms"]))
        if "x_mitre_aliases" in soft:
            row["aliases"] = ", ".join(sorted(soft["x_mitre_aliases"][1:]))
        row["type"] = soft["type"]  # malware or tool

        software_rows.append(row)

    citations = get_citations(software)
    dataframes = {
        "software": pd.DataFrame(software_rows).sort_values("name"),
    }
    # add relationships
    dataframes.update(relationshipsToDf(src, relatedType="software"))
    # add/merge citations
    if not citations.empty:
        if "citations" in dataframes:  # append to existing citations from references
            dataframes["citations"] = dataframes["citations"].append(citations)
        else:  # add citations
            dataframes["citations"] = citations

        dataframes["citations"].sort_values("reference")

    return dataframes

def groupsToDf(src, domain):
    """
    Parse STIX groups from the given data and return corresponding pandas dataframes.
    :param src: MemoryStore or other stix2 DataSource object holding the domain data
    :param domain: domain of ATT&CK src corresponds to, e.g "enterprise-attack"
    :returns: a lookup of labels (descriptors/names) to dataframes
    """

    groups = src.query([Filter("type", "=", "intrusion-set")])
    groups = remove_revoked_deprecated(groups)
    group_rows = []
    for group in tqdm(groups, desc="parsing groups"):
        row = parseBaseStix(group)
        # add group aliases
        if "aliases" in group:
            associated_groups = []
            associated_group_citations = []
            for alias in sorted(group["aliases"][1:]):
                # find the reference for the alias
                associated_groups.append(alias)
                for ref in group["external_references"]:
                    if ref["source_name"] == alias:
                        associated_group_citations.append(ref["description"])
                        break
                        # aliases.append(alias)
            row["associated groups"] = ", ".join(associated_groups)
            row["associated groups citations"] = ", ".join(associated_group_citations)

        group_rows.append(row)

    citations = get_citations(groups)
    dataframes = {
        "groups": pd.DataFrame(group_rows).sort_values("name"),
    }
    # add relationships
    dataframes.update(relationshipsToDf(src, relatedType="group"))
    # add/merge citations
    if not citations.empty:
        if "citations" in dataframes:  # append to existing citations from references
            dataframes["citations"] = dataframes["citations"].append(citations)
        else:  # add citations
            dataframes["citations"] = citations

        dataframes["citations"].sort_values("reference")

    return dataframes

def mitigationsToDf(src, domain):
    """
    Parse STIX mitigations from the given data and return corresponding pandas dataframes.
    :param src: MemoryStore or other stix2 DataSource object holding the domain data
    :param domain: domain of ATT&CK src corresponds to, e.g "enterprise-attack"
    :returns: a lookup of labels (descriptors/names) to dataframes
    """

    mitigations = src.query([Filter("type", "=", "course-of-action")])
    mitigations = remove_revoked_deprecated(mitigations)
    mitigation_rows = []
    for mitigation in tqdm(mitigations, desc="parsing mitigations"):
        mitigation_rows.append(parseBaseStix(mitigation))

    citations = get_citations(mitigations)
    dataframes = {
        "mitigations": pd.DataFrame(mitigation_rows).sort_values("name"),
    }
    # add relationships
    dataframes.update(relationshipsToDf(src, relatedType="mitigation"))
    # add/merge citations
    if not citations.empty:
        if "citations" in dataframes:  # append to existing citations from references
            dataframes["citations"] = dataframes["citations"].append(citations)
        else:  # add citations
            dataframes["citations"] = citations

        dataframes["citations"].sort_values("reference")

    return dataframes

class CellRange:
    """
    helper class for handling ranges of cells in a spreadsheet. Note: not 0-indexed, row and cols start at 1.
    Data is optional argument for data to store in the cellrange in the case of merged ranges
    format is a dict {name, format} for the XlsxWriter style. Formats of the same name will not be defined multiple
        times to the worksheet; only the first definition will be used
    """

    def __init__(self, leftCol, rightCol, topRow, bottomRow, data=None, format=None):
        self.leftCol = leftCol
        self.rightCol = rightCol
        self.topRow = topRow
        self.bottomRow = bottomRow
        self.data = data
        self.format = format

    def to_excel_format(self):
        """return the range in excel format, e.g A4:C7"""
        return f"{self._loc_to_excel(self.topRow, self.leftCol)}:{self._loc_to_excel(self.bottomRow, self.rightCol)}"

    def _loc_to_excel(self, row, col):
        """ Convert given row and column number to an Excel-style cell name. Note: not 0-indexed"""
        letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
        result = []
        while col:
            col, rem = divmod(col - 1, 26)
            result[:0] = letters[rem]
        return ''.join(result) + str(row)

def build_technique_and_sub_columns(src, techniques, columns, merge_data_handle, matrix_grid_handle, tactic_name):
    """
    Build technique and subtechnique columns for a given matrix and attach them to the appropriate object listings
    :param src: MemoryStore or other stix2 DataSource object holding the domain data
    :param techniques: List of technique stix objects belong in this tactic column
    :param columns: Existing columns in this matrix (used for placement)
    :param merge_data_handle: Handle to the 'merge' data object for this matrix
    :param matrix_grid_handle: Handle to the 2D grid array being constructed for the matrix (technique and subtechnique
                                columns will be appended here)
    :param tactic_name: The name of the corresponding tactic for this column
    :return: Nothing (meta - modifies the passed in merge_data_handle and matrix_grid_handle objects)
    """
    techniques_column = []
    subtechniques_column = []

    for technique in techniques:
        techniques_column.append(technique["name"])
        # sub-technique relationships
        subtechnique_ofs = src.query([
            Filter("type", "=", "relationship"),
            Filter("relationship_type", "=", "subtechnique-of"),
            Filter("target_ref", "=", technique["id"])
        ])
        if len(subtechnique_ofs) > 0:  # if there are sub-techniques on the tactic
            technique_top = len(techniques_column) + 1  # top of row range to merge
            # get sub-techniques
            subtechniques = [src.get(rel["source_ref"]) for rel in subtechnique_ofs]
            subtechniques = remove_revoked_deprecated(subtechniques)
            subtechniques = sorted(subtechniques, key=lambda x: x["name"])
            for i in range(len(subtechniques)):  # for each sub-technique
                if i != 0:
                    techniques_column.append("")  # first sub-technique is parallel to the technique in the layout
                subtechniques_column.append(subtechniques[i]["name"])
            technique_bottom = len(techniques_column) + 1  # bottom of row range to merge
            if technique_top != technique_bottom:  # more than 1 sub-technique
                merge_data_handle.append(CellRange(  # merge technique portion of cell group
                    len(columns),
                    len(columns),
                    technique_top,
                    technique_bottom,
                    data=technique["name"],
                    format={  # format of the merged range
                        "name": "supertechnique",
                        "format": {
                            'valign': 'vcenter',
                            'text_wrap': 1,
                            'shrink': 1
                        }
                    }
                ))
        else:  # no sub-techniques; add empty cell parallel to technique
            subtechniques_column.append("")
    # end adding techniques and sub-techniques to column

    matrix_grid_handle.append(techniques_column)  # add technique column to grid

    if len(list(filter(lambda x: x != "", subtechniques_column))) > 0:  # if there are sub-techniques for the tactic
        matrix_grid_handle.append(subtechniques_column)  # add sub-technique sub-column
        columns.append("")  # add empty tactic header for the sub-column
        merge_data_handle.append(  # merge tactic column header with the sub-column header that was just appended
            CellRange(len(columns) - 1,
                      len(columns),
                      1,
                      1,
                      data=tactic_name,
                      format={  # tactic header formatting
                          "name": "tacticHeader",
                          "format": {
                              "bold": 1,
                              "border": 1,
                              "font_size": 14,
                              "align": "center",
                              "shrink": 1
                          }
                      }
                      ))

def build_parsed_DF_matrix(matrix_grid, columns, merge, parsed_dict):
    """
    Build the DF matrix object
    :param matrix_grid: 2D array of the matrix to build
    :param columns: Column headers
    :param merge: Any applicable cell merge ranges and styles
    :param parsed_dict: Dictionary containing name and description for the matrix
    :return: { matrix, name, description, merge, border } where
        matrix is a pandas dataframe of the matrix
        name is the name of the matrix
        description is the description of the matrix
        merge is a list of CellRange objects that need to be merged for formatting of the sub-techniques in the matrix
        columns is the number of columns in the data
    """
    parsed = copy.deepcopy(parsed_dict)
    # reshape array so that pandas consumes it properly
    matrix_grid = np.flip(np.rot90(matrix_grid), 0)
    # create dataframe for array
    df = pd.DataFrame(matrix_grid, columns=columns)

    # Set additional data for the matrix
    parsed["matrix"] = df  # actual dataframe
    parsed["merge"] = merge  # merge ranges and associated styles
    parsed["columns"] = len(columns)  # number of columns with data
    return parsed

def matricesToDf(src, domain):
    """
    Parse STIX matrices from the given data and return parsed matrix structures
    :param src: MemoryStore or other stix2 DataSource object holding the domain data
    :param domain: domain of ATT&CK src corresponds to, e.g "enterprise-attack"
    :returns: [{ matrix, name, description, merge, border }, ... ] where
        matrix is a pandas dataframe of the matrix
        name is the name of the matrix
        description is the description of the matrix
        merge is a list of CellRange objects that need to be merged for formatting of the sub-techniques in the matrix
        columns is the number of columns in the data
    """
    matrices = src.query([Filter("type", "=", "x-mitre-matrix")])
    matrices = remove_revoked_deprecated(matrices)
    matrices_parsed = []
    sub_matrices_parsed = []

    for matrix in tqdm(matrices, desc="parsing matrices"):
        sub_matrices_grid = dict()
        sub_matrices_merges = dict()
        sub_matrices_columns = dict()
        for entry in MATRIX_PLATFORMS_LOOKUP[domain]:
            sub_matrices_grid[entry] = []
            sub_matrices_merges[entry] = []
            sub_matrices_columns[entry] = []

        parsed = {
            "name": matrix['name'] if len(matrices) == 1 else F"{domain.split('-')[0].capitalize()} {matrix['name']}",
            "description": matrix["description"]
        }

        matrix_grid = []  # matrix layout in 2d array
        merge = []  # list of CellRange objects to merge later

        columns = []  # column names
        for tactic_ref in tqdm(matrix["tactic_refs"], desc="processing matrix tactics"):
            tactic = src.get(tactic_ref)
            columns.append(tactic["name"])  # add tactic header

            # parse techniques in tactic
            techniques = list(
                filter(lambda t: not ("x_mitre_is_subtechnique" in t and t["x_mitre_is_subtechnique"]), src.query([
                    Filter("type", "=", "attack-pattern"),
                    Filter("kill_chain_phases.phase_name", "=", tactic["x_mitre_shortname"]),
                ])))
            techniques = remove_revoked_deprecated(techniques)
            techniques = sorted(techniques, key=lambda x: x["name"])
            # add techniques
            build_technique_and_sub_columns(src, techniques, columns, merge, matrix_grid, tactic['name'])

            for platform in MATRIX_PLATFORMS_LOOKUP[domain]:
                # In order to support "groups" of platforms, each platform is checked against the lookup a second time.
                # If an second entry can be found, the results from that query will be used, otherwise, the singular
                # platform will be.
                a_techs = filter_platforms(techniques,
                                           [platform] if platform not in MATRIX_PLATFORMS_LOOKUP
                                           else MATRIX_PLATFORMS_LOOKUP[platform])
                if a_techs:
                    sub_matrices_columns[platform].append(tactic['name'])
                    build_technique_and_sub_columns(src, a_techs, sub_matrices_columns[platform],
                                                    sub_matrices_merges[platform], sub_matrices_grid[platform],
                                                    tactic['name'])

        # square the grid because pandas doesn't like jagged columns
        longest_column = 0
        for column in matrix_grid:
            longest_column = max(len(column), longest_column)
        for column in matrix_grid:
            for i in range((longest_column - len(column))):
                column.append("")

        for submatrix in sub_matrices_grid:
            mg = sub_matrices_grid[submatrix]
            for column in mg:
                longest_column = max(len(column), longest_column)
            for column in mg:
                for i in range((longest_column - len(column))):
                    column.append("")
        # matrix is now squared

        parsed = build_parsed_DF_matrix(matrix_grid, columns, merge, parsed)
        matrices_parsed.append(parsed)

        for submatrix in sub_matrices_grid:
            if sub_matrices_grid[submatrix]:  # make sure we found matches for something
                local = copy.deepcopy(parsed)
                local['name'] = f"{submatrix}" if len(matrices) == 1 else f"{submatrix} {matrix['name']}"
                local['description'] = local['description'].split(":")[0] + f": {submatrix}"
                subparsed = build_parsed_DF_matrix(sub_matrices_grid[submatrix], sub_matrices_columns[submatrix],
                                                   sub_matrices_merges[submatrix], local)
                sub_matrices_parsed.append(subparsed)

    # end adding of matrices
    return matrices_parsed, sub_matrices_parsed

def relationshipsToDf(src, relatedType=None):
    """
    Parse STIX relationships from the given data and return corresponding pandas dataframes.
    :param src: MemoryStore or other stix2 DataSource object holding the domain data
    :param relatedType: optional, singular attack type to only return relationships with, e.g "mitigation"
    :returns: a lookup of labels (descriptors/names) to dataframes
    """

    # Helper lookups
    attackToStixTerm = {
        "technique": ["attack-pattern"],
        "tactic": ["x-mitre-tactic"],
        "software": ["tool", "malware"],
        "group": ["intrusion-set"],
        "mitigation": ["course-of-action"],
        "matrix": ["x-mitre-matrix"],
    }
    stixToAttackTerm = {
        "attack-pattern": "technique",
        "x-mitre-tactic": "tactic",
        "tool": "software",
        "malware": "software",
        "intrusion-set": "group",
        "course-of-action": "mitigation",
        "x-mitre-matrix": "matrix"
    }

    # get master list of relationships
    relationships = src.query([Filter("type", "=", "relationship")])
    relationships = remove_revoked_deprecated(relationships)
    relationship_rows = []  # build list of rows for dataframe
    # tqdm description depends on the related type and parameters
    iterdesc = "parsing all relationships" if not relatedType else f"parsing relationships for type={relatedType}"
    for relationship in tqdm(relationships, desc=iterdesc):
        source = src.get(relationship["source_ref"])  # source object of the relationship
        target = src.get(relationship["target_ref"])  # target object of the relationship

        # filter if related objects don't exist or are revoked or deprecated
        if not source or source.get("x_mitre_deprecated", False) is True or source.get("revoked", False) is True:
            continue
        if not target or target.get("x_mitre_deprecated", False) is True or target.get("revoked", False) is True:
            continue
        if relationship["relationship_type"] == "revoked":
            continue

        # don't track sub-technique relationships, those are tracked in the techniques df
        if relationship["relationship_type"] == "subtechnique-of":
            continue

        # filter out relationships not with relatedType
        if relatedType:
            related = False
            for stixTerm in attackToStixTerm[relatedType]:  # try all stix types for the ATT&CK type
                if (source["type"] == stixTerm or
                        target["type"] == stixTerm):  # if any stix type is part of the relationship
                    related = True
                    break
            if not related: continue  # skip this relationship if the types don't match

        # add mapping data
        row = {}

        def add_side(label, sdo):
            """add data for one side of the mapping"""
            if "external_references" in sdo and sdo["external_references"][0]["source_name"] in ["mitre-attack",
                                                                                "mitre-mobile-attack"]:
                row[f"{label} ID"] = sdo["external_references"][0]["external_id"]  # "source ID" or "target ID"
            if "name" in sdo:
                row[f"{label} name"] = sdo["name"]  # "source name" or "target name"
            row[f"{label} type"] = stixToAttackTerm[sdo["type"]]  # "source type" or "target type"

        add_side("source", source)
        row["mapping type"] = relationship["relationship_type"]  # mapping type goes between the source/target data
        add_side("target", target)
        if "description" in relationship:  # add description of relationship to the end of the row
            row["mapping description"] = relationship["description"]

        relationship_rows.append(row)

    citations = get_citations(relationships)
    relationships = pd.DataFrame(relationship_rows).sort_values(
        ["mapping type", "source type", "target type", "source name", "target name"])

    if not relatedType:  # return all relationships and citrations
        dataframes = {
            "relationships": relationships,
        }
        if not citations.empty:
            dataframes["citations"] = citations.sort_values("reference")

        return dataframes
    else:  # break into dataframes by mapping type
        dataframes = {}

        # group:software / "associated {other type}"
        relatedGroupSoftware = relationships.query(
            "`mapping type` == 'uses' and (`source type` == 'group' or `source type` == 'software') and "
            "(`target type` == 'group' or `target type` == 'software')")
        if not relatedGroupSoftware.empty:
            dataframes[f"associated {'software' if relatedType == 'group' else 'groups'}"] = relatedGroupSoftware

        # technique:group + technique:software / "procedure examples"
        procedureExamples = relationships.query("`mapping type` == 'uses' and `target type` == 'technique'")
        if not procedureExamples.empty:
            dataframes["procedure examples" if relatedType == "technique" else "techniques used"] = procedureExamples

        # technique:mitigation / "mitigation mappings"
        relatedMitigations = relationships.query("`mapping type` == 'mitigates'")
        if not relatedMitigations.empty:
            dataframes[
                'associated mitigations' if relatedType == 'technique' else 'techniques addressed'] = relatedMitigations

        if not citations.empty:
            # fitler citations by ones actually used
            # build master list of used citations
            usedCitations = set()
            for dfname in dataframes:
                df = dataframes[dfname]
                for description in filter(lambda x: x == x, df[
                    "mapping description"].tolist()):  # filter out missing descriptions which for whatever reason
                    # in pandas don't equal themselves
                    [usedCitations.add(x) for x in re.findall(r"\(Citation: (.*?)\)", description)]

            citations = citations[citations.reference.isin(list(usedCitations))]  # filter to only used references

            dataframes["citations"] = citations.sort_values("reference")

        return dataframes


import ROOT
import math

def heater_test_function( x, p ):
  # x[0] x value (time axis)
  # p[0] Delta t1 amplitude of temp jump
  # p[1] tau1     rising time constant (s)
  # p[2] m        background slope (K/s) added relative to turn on time
  # p[3] b        background temperature (K)
  # p[4] t1       turn on time (s)
  # p[5] t2       turn off time (s)
  # p[6] Delta t2 amplitude of temp fall
  # p[7] tau2     falling time constant (s)

  if x[0] < p[4]:  # before heater turns on
    return p[3] + p[2]*(x[0]-p[4]);
  elif x[0] < p[5]: #while heater is on
    retval = p[3] + p[2]*(x[0]-p[4]);
    retval = retval + p[0]*(1 - math.exp( -(x[0]-p[4])/p[1] ) )
    return retval
  else:
    # heater turns off
    bprime = p[3] + p[0]*( 1 - math.exp( -(p[5]-p[4])/p[1] ) ) - p[6]
    retval = bprime + p[2]*(x[0]-p[4]);
    retval = retval + p[6]*math.exp( -(x[0]-p[5])/p[7] )
    return retval

def htf_sameEXP( x, p ): #heater test function with same exponetial values
  # x[0] x value (time axis)
  # p[0] Delta t1 amplitude of temp jump
  # p[1] tau1     rising time constant (s)
  # p[2] m        background slope (K/s) added relative to turn on time
  # p[3] b        background temperature (K)
  # p[4] t1       turn on time (s)
  # p[5] t2       turn off time (s)
  # p[6] Delta t2 amlitude of temp fall

  if x[0] < p[4]:  # before heater turns on
    return p[3] + p[2]*(x[0]-p[4]);
  elif x[0] < p[5]: #while heater is on
    retval = p[3] + p[2]*(x[0]-p[4]);
    retval = retval + p[0]*(1 - math.exp( -(x[0]-p[4])/p[1] ) )
    return retval
  else:
    # heater turns off
    bprime = p[3] + p[0]*( 1 - math.exp( -(p[5]-p[4])/p[1] ) ) - p[6]
    retval = bprime + p[2]*(x[0]-p[4]);
    retval = retval + p[6]*math.exp( -(x[0]-p[5])/p[1] )
    return retval

def htf_sameDT( x, p ): #heater test function with same exponetial values
  # x[0] x value (time axis)
  # p[0] Delta t1 amplitude of temp jump
  # p[1] tau1     rising time constant (s)
  # p[2] m        background slope (K/s) added relative to turn on time
  # p[3] b        background temperature (K)
  # p[4] t1       turn on time (s)
  # p[5] t2       turn off time (s)
  #               Delta t2 amlitude of temp fall
  # p[6] tau2
  if x[0] < p[4]:  # before heater turns on
    return p[3] + p[2]*(x[0]-p[4]);
  elif x[0] < p[5]: #while heater is on
    retval = p[3] + p[2]*(x[0]-p[4]);
    retval = retval + p[0]*(1 - math.exp( -(x[0]-p[4])/p[1] ) )
    return retval
  else:
    # heater turns off
    bprime = p[3] + p[0]*( 1 - math.exp( -(p[5]-p[4])/p[1] ) ) - p[0]
    retval = bprime + p[2]*(x[0]-p[4]);
    retval = retval + p[0]*math.exp( -(x[0]-p[5])/p[6] )
    return retval

def htf_sameTexp( x, p ): #heater test function with same exponetial values
  # x[0] x value (time axis)
  # p[0] Delta t1 amplitude of temp jump
  # p[1] tau1     rising time constant (s)
  # p[2] m        background slope (K/s) added relative to turn on time
  # p[3] b        background temperature (K)
  # p[4] t1       turn on time (s)
  # p[5] t2       turn off time (s)
  #               Delta t2 amlitude of temp fall
  #               tau2=tau1

  if x[0] < p[4]:  # before heater turns on
    return p[3] + p[2]*(x[0]-p[4]);
  elif x[0] < p[5]: #while heater is on
    retval = p[3] + p[2]*(x[0]-p[4]);
    retval = retval + p[0]*(1 - math.exp( -(x[0]-p[4])/p[1] ) )
    return retval
  else:
    # heater turns off
    bprime = p[3] + p[0]*( 1 - math.exp( -(p[5]-p[4])/p[1] ) ) - p[0]
    retval = bprime + p[2]*(x[0]-p[4]);
    retval = retval + p[0]*math.exp( -(x[0]-p[5])/p[1] )
    return retval

def cool_function(x, p): #cooling portion only
    # x[0] x value (time axis)
    # p[0] Delta t1 amplitude of temp fall
    # p[1] tau1     rising time constant (s)
    # p[2] m        background slope (K/s) added relative to turn on time
    # p[3] b        background temperature (K)
    # p[4] t1       turn off time (s)

    if x[0] < p[4]: # while heater turns on
        return p[3] + p[2]*(x[0]-p[4])+p[0];
    else:
        # heater turns off
        retval = p[3] + p[2]*(x[0]-p[4]);
        retval = retval + p[0]*math.exp( -(x[0]-p[4])/p[1] )
        return retval

def high_function( x, p ):
  # x[0] x value (time axis)
  # p[0] Delta t1 amplitude of temp jump
  # p[1] tau1     rising time constant (s)
  #### p[2] m        background slope (K/s) added relative to turn on time
  # p[2] b        background temperature (K)
  # p[3] t1       turn on time (s)

  if x[0] < p[3]:  # before heater turns on
    return p[2]# + p[2]*(x[0]-p[4])
  else: #while heater is on
    retval = p[2] #+ p[2]*(x[0]-p[4]);
    retval = retval + p[0]*(1 - math.exp( -(x[0]-p[3])/p[1] ) )
    return retval

def linear_function( x, p ): #heater test function with same exponetial values
  # x[0] x value (time axis)
  # p[0] m        background slope (K/s) added relative to turn on time
  # p[1] b        background temperature (K)
  return p[1] + p[0]*(x[0]);

def SingleExpoWithSlopeBackground(oS, pp0 ,pp1 , pp2, pp3):
    SingleExpoWithSlopeBackground = ROOT.TF1('SingleExpoWithBackground', '[0]*exp((-x+{0})/[1]) + [2]*(x-{0})+[3]'.format(oS))
    SingleExpoWithSlopeBackground.SetParameters(pp0,pp1,pp2,pp3 )
    SingleExpoWithSlopeBackground.SetParName(0, '#Delta Temp')
    SingleExpoWithSlopeBackground.SetParName(1, 'time constant #tau')
    SingleExpoWithSlopeBackground.SetParName(2, 'Background Slope')
    SingleExpoWithSlopeBackground.SetParName(3, 'Final Temp')
    SingleExpoWithSlopeBackground.SetParLimits(0, -1, 1)
    SingleExpoWithSlopeBackground.SetParLimits(1, -1000, 1000)
    SingleExpoWithSlopeBackground.SetParLimits(2, -1, 1)
    SingleExpoWithSlopeBackground.SetParLimits(3, 0, 2)
    return SingleExpoWithSlopeBackground

def HeaterTestFunction(pars, tmin, tmax):
    HeaterTestFunction = ROOT.TF1('HeaterTest', heater_test_function, tmin, tmax, 8 )
    HeaterTestFunction.SetParameters( pars[0], pars[1], pars[2], pars[3], pars[4], pars[5], pars[6], pars[7] )
    HeaterTestFunction.SetParName(0, '#Delta T_{1} (K)')
    HeaterTestFunction.SetParName(1, '#tau_{1} (s)    ')
    HeaterTestFunction.SetParName(2, 'BG slope (K/s)  ')
    HeaterTestFunction.SetParName(3, 'Offset T        ')
    HeaterTestFunction.SetParName(4, 'turn on (s)     ')
    HeaterTestFunction.SetParName(5, 'turn off (s)    ')
    HeaterTestFunction.SetParName(6, '#Delta T_{2} (K)')
    HeaterTestFunction.SetParName(7, '#tau_{2} (s)    ')
    HeaterTestFunction.SetParLimits(0, 0, 1.0)
    HeaterTestFunction.SetParLimits(1, 1, 1000)
    HeaterTestFunction.SetParLimits(2, -0.1, 0.1)
    HeaterTestFunction.SetParLimits(3, 0, 2)
    HeaterTestFunction.SetParLimits(4, 0, 14000)
    HeaterTestFunction.SetParLimits(5, 0, 14000)
    HeaterTestFunction.SetParLimits(6, 0, 1.0)
    HeaterTestFunction.SetParLimits(7, 1, 1000)
    HeaterTestFunction.SetNpx(300)
    return HeaterTestFunction

def HTF_Same_exp(pars, tmin, tmax):
    HTF_Same_exp = ROOT.TF1('HeaterTest', htf_sameEXP, tmin, tmax, 7 )
    HTF_Same_exp.SetParameters( pars[0], pars[1], pars[2], pars[3], pars[4], pars[5], pars[6] )

    HTF_Same_exp.SetParName(0, '#Delta T_{1} (K)')
    HTF_Same_exp.SetParName(1, '#tau_{1} (s)    ')
    HTF_Same_exp.SetParName(2, 'BG slope (K/s)  ')
    HTF_Same_exp.SetParName(3, 'Offset T        ')
    HTF_Same_exp.SetParName(4, 'turn on (s)     ')
    HTF_Same_exp.SetParName(5, 'turn off (s)    ')
    HTF_Same_exp.SetParName(6, '#Delta T_{2} (K)')
    HTF_Same_exp.SetParLimits(0, 0, 1.0)
    HTF_Same_exp.SetParLimits(1, 1, 1000)
    HTF_Same_exp.SetParLimits(2, -0.1, 0.1)
    HTF_Same_exp.SetParLimits(3, 0, 2)
    HTF_Same_exp.SetParLimits(4, 0, 14000)
    HTF_Same_exp.SetParLimits(5, 0, 14000)
    HTF_Same_exp.SetParLimits(6, 0, 1.0)
    HTF_Same_exp.SetNpx(300)
    return HTF_Same_exp

def HTF_Same_T(pars, tmin, tmax):
    HTF_Same_exp = ROOT.TF1('HeaterTest', htf_sameDT, tmin, tmax, 7 )
    HTF_Same_exp.SetParameters( pars[0], pars[1], pars[2], pars[3], pars[4], pars[5], pars[6] )

    HTF_Same_exp.SetParName(0, '#Delta T_{1} (K)')
    HTF_Same_exp.SetParName(1, '#tau_{1} (s)    ')
    HTF_Same_exp.SetParName(2, 'BG slope (K/s)  ')
    HTF_Same_exp.SetParName(3, 'Offset T        ')
    HTF_Same_exp.SetParName(4, 'turn on (s)     ')
    HTF_Same_exp.SetParName(5, 'turn off (s)    ')
    HTF_Same_exp.SetParName(6, '#tau_{2} (s)    ')
    HTF_Same_exp.SetParLimits(0, 0, 1.0)
    HTF_Same_exp.SetParLimits(1, 1, 1000)
    HTF_Same_exp.SetParLimits(2, -0.1, 0.1)
    HTF_Same_exp.SetParLimits(3, 0, 2)
    HTF_Same_exp.SetParLimits(4, 0, 14000)
    HTF_Same_exp.SetParLimits(5, 0, 14000)
    HTF_Same_exp.SetParLimits(6, 0, 1000)
    HTF_Same_exp.SetNpx(300)
    return HTF_Same_exp

def HTF_Same_Texp(pars, tmin, tmax):
    HTF_Same_exp = ROOT.TF1('HeaterTest', htf_sameTexp, tmin, tmax, 6 )
    HTF_Same_exp.SetParameters( pars[0], pars[1], pars[2], pars[3], pars[4], pars[5])

    HTF_Same_exp.SetParName(0, '#Delta T_{1} (K)')
    HTF_Same_exp.SetParName(1, '#tau_{1} (s)    ')
    HTF_Same_exp.SetParName(2, 'BG slope (K/s)  ')
    HTF_Same_exp.SetParName(3, 'Offset T        ')
    HTF_Same_exp.SetParName(4, 'turn on (s)     ')
    HTF_Same_exp.SetParName(5, 'turn off (s)    ')
    HTF_Same_exp.SetParLimits(0, 0, 1.0)
    HTF_Same_exp.SetParLimits(1, 1, 1000)
    HTF_Same_exp.SetParLimits(2, -0.1, 0.1)
    HTF_Same_exp.SetParLimits(3, 0, 2)
    HTF_Same_exp.SetParLimits(4, 0, 14000)
    HTF_Same_exp.SetParLimits(5, 0, 14000)
    HTF_Same_exp.SetNpx(300)
    return HTF_Same_exp

def Cool_Test(pars, tmin, tmax):
    Cool_Test = ROOT.TF1('HeaterTest', cool_function, tmin, tmax, 5 )
    Cool_Test.SetParameters( pars[0], pars[1], pars[2], pars[3], pars[4] )
    Cool_Test.SetParName(0, '#Delta T_{2} (K)')
    Cool_Test.SetParName(1, '#tau_{2} (s)    ')
    Cool_Test.SetParName(2, 'BG slope (K/s)  ')
    Cool_Test.SetParName(3, 'Offset T        ')
    Cool_Test.SetParName(4, 'turn off (s)    ')
    Cool_Test.SetParLimits(0, 0, 1.0)
    Cool_Test.SetParLimits(1, 1, 1000)
    Cool_Test.SetParLimits(2, -0.1, 0.1)
    Cool_Test.SetParLimits(3, 0, 2)
    Cool_Test.SetParLimits(4, 0, 14000)
    Cool_Test.SetNpx(300)
    return Cool_Test

def High_Test(pars, tmin, tmax):
    High_Test = ROOT.TF1('HeaterTest', high_function, tmin, tmax, 4 )
    High_Test.SetParameters( pars[0], pars[1], pars[2], pars[3])
    High_Test.SetParName(0, '#Delta T_{1} (K) ')
    High_Test.SetParName(1, '#tau_{1} (s)     ')
    #High_Test.SetParName(2, 'BG slope (K/s)   ')
    High_Test.SetParName(2, 'Starting Temp (K)')
    High_Test.SetParName(3, 'turn on (s)      ')
    High_Test.SetParLimits(0, 0, 1.0)
    High_Test.SetParLimits(1, 1, 1000)
    #High_Test.SetParLimits(2, -0.1, 0.1)
    High_Test.SetParLimits(2, 0., 2.)
    High_Test.SetParLimits(3, -10, 10)
    High_Test.SetNpx(300)
    return High_Test

def Linear_Function(pars, tmin, tmax):
    Linear_Test = ROOT.TF1('HeaterTest', linear_function, tmin, tmax, 2 )
    Linear_Test.SetParameters( pars[0], pars[1] )
    Linear_Test.SetParName(0, 'BG slope (K/s)  ')
    Linear_Test.SetParName(1, 'Offset T        ')
    Linear_Test.SetParLimits(0, -0.1, 0.1)
    Linear_Test.SetParLimits(1, 0, 2)
    return Linear_Test

#test1 = HeaterTestFunction( [0.024, 32.0, 5.0e-6, 1.16, 900.0, 1000.0, 0.024, 32.0 ], 600.0, 1200.0 )

#tc = ROOT.TCanvas()
#test1.Draw()
#tc.Print("plot_heater_test_function.png")


"""
Created on Sun Feb  2 13:28:48 2020

@author: matias
"""
import numpy as np
from numpy.linalg import inv
from matplotlib import pyplot as plt
import time
#import camb #No lo reconoce la compu del df
from scipy.integrate import cumtrapz as cumtrapz
from scipy.integrate import simps as simps
from scipy.interpolate import interp1d
from scipy.constants import c as c_luz #metros/segundos
c_luz_km = c_luz/1000
#%%

def E_LCDM(z, omega_m):
    omega_lambda = 1 - omega_m
    E = np.sqrt(omega_m * (1 + z)**3 + omega_lambda)
    return E

def H_LCDM(z, omega_m, H_0):
    H = H_0 * E_LCDM(z,omega_m)
    return H

def H_LCDM_rad(z, omega_m, H_0):
    omega_r = 4.18343*10**(-5) / (H_0/100)**2
    omega_lambda = 1 - omega_m - omega_r

    if isinstance(z, (np.ndarray, list)):
        H = H_0 * np.sqrt(omega_r * (1 + z)**4 + omega_m * (1 + z)**3 + omega_lambda)
    else:
        H = H_0 * (omega_r * (1 + z)**4 + omega_m * (1 + z)**3 + omega_lambda)**(1/2)

    return H

#%%

if __name__ == '__main__':
    from matplotlib import pyplot as plt
    import sys
    import os
    from os.path import join as osjoin
    from pc_path import definir_path
    path_git, path_datos_global = definir_path()
    os.chdir(path_git)
    sys.path.append('./Software/Funcionales/')
    from funciones_data import leer_data_BAO,leer_data_AGN

    #%% BAO
    os.chdir(path_git+'/Software/Estadística/Datos/BAO/')
#    os.chdir(path_git+'/Software/Estadística/Datos/BAO/Datos_sin_nuevos')
    dataset_BAO = []
    archivo_BAO = ['datos_BAO_da.txt','datos_BAO_dh.txt','datos_BAO_dm.txt',
                    'datos_BAO_dv.txt','datos_BAO_H.txt']
    for i in range(5):
        aux = leer_data_BAO(archivo_BAO[i])
        dataset_BAO.append(aux)

    [omega_m,H_0] = [0.311, 66.012] #Usando los datos nuevos
    #[omega_m,H_0] = [0.314, 66.432] #Usando datos nuevos (rd camb)

    #[omega_m,H_0] = [0.302, 62.806] #Sin usar datos nuevos
    #[omega_m,H_0] = [0.301, 66.003] #Sin usar datos nuevos (rd camb)

    theta = [omega_m, H_0]

    (z_data,valores,errores_cuad,wd_fid)=dataset_BAO[i]

    zs = np.linspace(0.01, 3, 10**6)
    H_teo = H_LCDM(zs, omega_m, H_0)
    z_data
    rd = r_drag(omega_m,H_0,wd_fid[0]) #Calculo del rd

    #%matplotlib qt5
    legends = ['Da_rd','Dh_rd','Dm_rd','Dv_rd','H*rd']
    num_datasets=5
    chies_BAO = np.zeros(num_datasets)
    for i in range(5):
        (z_data_BAO, valores_data, errores_data_cuad,wb_fid) = dataset_BAO[i]
        if i==0: #Dato de Da
            #La parte teo del chi2
            rd = r_drag_camb(omega_m,H_0,wb_fid) #Calculo del rd
            distancias_teoricas = Hs_to_Ds(zs,H_teo,z_data_BAO,i) #Bien hecho!
            output_teorico = Ds_to_obs_final(zs, distancias_teoricas, rd, i)
        else: #De lo contrario..
            distancias_teoricas = Hs_to_Ds(zs,H_teo,z_data_BAO,i)
            output_teorico = np.zeros(len(z_data_BAO))
            for j in range(len(z_data_BAO)):
                #La parte teo del chi2
                 rd = r_drag_camb(omega_m,H_0,wb_fid[j]) #Calculo del rd
                 output_teorico[j] = Ds_to_obs_final(zs,distancias_teoricas[j],rd,i) #Bien hecho!
        #Calculo el chi2 para cada tipo de dato (i)
        chies_BAO[i] = chi_2_BAO(output_teorico,valores_data,errores_data_cuad)
    print(np.sum(chies_BAO)/(20-2))
#%%
    plt.figure(2)
    plt.grid()
    plt.errorbar(z_data_BAO,valores_data,np.sqrt(errores_data_cuad),fmt='.r')
    plt.plot(z_data_BAO,output_teorico,'.-b')#los puntos unidos por lineas, no es la forma real
    plt.title(legends[i])
    plt.xlabel('z (redshift)')
    plt.ylabel(legends[i])
    os.chdir(path_git+'/Software/Estadística/Datos/BAO/Imagen')
    #os.chdir(path_git+'/Software/Estadística/Datos/BAO/Datos_sin_nuevos/Imagen')
    plt.savefig('{}'.format(legends[i]))
    plt.close()
    #Hay algo mal, los chi2 no coinciden con los grafs. Revisar los errores_data_cuad.

#%%
    #Data AGN
    os.chdir(path_git+'/Software/Estadística/Datos/Datos_AGN')
    data_agn = leer_data_AGN('table3.dat')
    H_0 = 73.48
    params_fijos = H_0
    omegas = np.linspace(0.1,1,50)
    output = np.zeros(len(omegas))
    for i,omega in enumerate(omegas):
        theta = omega
        output[i] = params_to_chi2_AGN(theta, params_fijos, data_agn)
    #%matplotlib qt5
    plt.plot(omegas,output/(len(data_agn[0])-1))
    plt.title('Chi2 reducido vs omega_m')
    index = np.where(output==min(output))[0][0]
    omegas[index]
    min(output/(len(data_agn[0])-1))

    #%%
    #Defino los parámetros que voy a utilizar
    omega_m = 0.3
    #H_0 = params_fijos
    gamma = 0.648
    beta = 7.735
    egamma = 0.007
    ebeta = 0.244
    #Importo los datos
    data_agn = leer_data_AGN('table3.dat')
    z_data, logFuv, eFuv, logFx, eFx  = data_agn

    zs = np.linspace(0,30,100000)
    Es = E_LCDM(zs, omega_m)
    INT = cumtrapz(Es**(-1), zs, initial=0)
    DlH0 = (c_luz_km * (1 + zs)) * INT
    output = interp1d(zs,DlH0)
    DlH0_teo = np.log10(output(z_data))
    #-np.log10(4*np.pi)/2
    DlH0_obs =  np.log10(3.24)-25+(logFx - gamma * logFuv - beta) / (2*gamma - 2)
    df_dgamma = (1/(2*(gamma-1))**2) * (-logFx + beta + logFuv)
    eDlH0_cuad = (eFx**2 + gamma**2 * eFuv**2 + ebeta**2) / (2*gamma - 2)**2 + (df_dgamma)**2 * egamma**2 #El cuadrado de los errores
    chi2_AGN = chi_2_AGN(DlH0_teo, DlH0_obs, eDlH0_cuad)

    #%matplotlib qt5
    plt.figure()
    plt.xlabel('z (redshift)')
    plt.ylabel(r'$\log(d_{L}H_{0})$')
    plt.errorbar(z_data,DlH0_obs,np.sqrt(eDlH0_cuad),marker='.',linestyle='')
    plt.plot(z_data,DlH0_obs,'.r')
    plt.plot(z_data,DlH0_teo)
    print(chi2_AGN/(len(data_agn[0])-1))

    N = len(data_agn[0]) #Número de datos
    P = 1 #Número de parámetros
    np.sqrt(2*(N-P))/(N-P) #Fórmula del error en el chi2 reducido

    #%%
    H_0 = 70
    omega_m = 0.83
    gamma = 0.6
    beta = 8.3
    delta = 0.63
    theta = [omega_m,beta,gamma,delta]

    #params_to_chi2_AGN_nuisance(theta, _, data_agn)/(len(z_data)-4)

    data_agn = leer_data_AGN('table3.dat')
    z_data_1, logFuv_1, eFuv_1, logFx_1, eFx_1  = data_agn

    zmin = 0
    zmax = 100
    mask = (z_data_1 > zmin) & (z_data_1 < zmax)

    z_data = z_data_1[mask]
    logFuv = logFuv_1[mask]
    logFx = logFx_1[mask]
    eFx = eFx_1[mask]
    eFuv = eFuv_1[mask]

    zs_modelo = np.linspace(0,30,10**6)
    Dl_teo = -np.log10(H_0) + zs_2_logDlH0(zs_modelo,omega_m,z_data)
    Dl_teo_cm = Dl_teo - np.log10(3.24) + 25
    psi = beta + gamma * logFuv + 2 * (gamma-1) * (Dl_teo_cm + 0.5 * np.log10(4*np.pi))

    #si_2 = eFx**2 + (gamma * eFuv)**2 + np.exp(2*np.log(delta)) #El cuadrado de los errores
    si_2 = eFx**2 + (gamma * eFuv)**2 + delta**2 #El cuadrado de los errores

    chi2_AGN = np.sum( ((logFx-psi)**2/si_2) + np.log(2*np.pi*si_2)) # menos en el paper

    print(chi2_AGN)
    print(chi2_AGN/(len(z_data)-3))

    plt.figure()
    plt.xlabel('z (redshift)')
    plt.ylabel(r'$Fx$')
    plt.errorbar(z_data,psi,np.sqrt(si_2),marker='.',linestyle='')
    plt.plot(z_data,logFx,'.r')


# -*- coding: utf-8 -*-
"""
  Author: kun.wang
  Create: 2015-06-11

using: PySide
"""

__author__ = 'kun.wang'

import sys

import PySide.QtCore as QtCore
import PySide.QtGui as QtGui

import honey
import honey_context as hc

class AddProjectDialog(QtGui.QDialog):
    def __init__(self, parent=None):
        super(AddProjectDialog, self).__init__(parent)
        self.setWindowTitle(u'添加项目')
        self.resize(400, 300)

        self.project_name = QtGui.QLineEdit()
        self.project_detail = QtGui.QTextEdit()

        self.ok = QtGui.QPushButton(u"创建")
        self.cancel = QtGui.QPushButton(u"取消")

        self.ok.clicked.connect(self.accept)
        self.cancel.clicked.connect(self.reject)

        self.__layout = QtGui.QGridLayout()
        self.__layout.setColumnStretch(1, 1)
        self.__layout.setColumnStretch(2, 1)
        self.__layout.addWidget(QtGui.QLabel(u'项目名称: '), 0, 0, QtCore.Qt.AlignRight)
        self.__layout.addWidget(self.project_name, 0, 1, 1, 2)
        self.__layout.addWidget(QtGui.QLabel(u'项目描述: '), 1, 0, QtCore.Qt.AlignRight)
        self.__layout.addWidget(self.project_detail, 1, 1, 3, 2)
        self.__layout.addWidget(self.ok, 4, 1)
        self.__layout.addWidget(self.cancel, 4, 2)

        self.setLayout(self.__layout)

    def get_data(self):
        data = dict()
        data["project_name"] = unicode(self.project_name.text())
        data["project_detail"] = unicode(self.project_detail.toPlainText())
        return data

class ProjectItem(QtGui.QWidget):
    bg_color = QtGui.QColor(150, 150, 250)
    focus_color = QtGui.QColor(100, 100, 130, 100)
    name_font = QtGui.QFont()

    def __init__(self, project, parent=None):
        super(ProjectItem, self).__init__(parent)
        self.project = project
        self.on_focus = False

    def paintEvent(self, event):
        painter = QtGui.QPainter(self)
        # draw background
        painter.fillRect(self.rect(), ProjectItem.bg_color)
        painter.fillRect(0, 0, 20, self.height(), QtGui.QColor(100, 100, 100))

        if self.on_focus:
            painter.fillRect(self.rect(), ProjectItem.focus_color)

        painter.setPen(QtGui.QPen(QtGui.QColor(255, 255, 255)))
        ProjectItem.name_font.setPixelSize(30)
        painter.setFont(ProjectItem.name_font)
        painter.drawText(30, self.height() - 10, "%s" % self.project.name())

    def mouseDoubleClickEvent(self, event):
        hc.Project = self.project
        # self.doubleClicked.emit(self)
        hc.Context.go_to_project.emit()

    def enterEvent(self, event):
        self.on_focus = True
        self.repaint()

    def leaveEvent(self, event):
        self.on_focus = False
        self.repaint()

class ProjectListView(QtGui.QWidget):
    def __init__(self, parent=None):
        super(ProjectListView, self).__init__(parent)
        self.item_height = 100
        self.item_space = 5
        self.items = []

        self.setMinimumWidth(200)
        self.setMinimumHeight(self.item_height + self.item_space)

    def add_item(self, item):
        item.setParent(self)
        item.resize(self.width(), self.item_height)
        self.items.append(item)
        item.show()
        height = len(self.items) * (self.item_height + self.item_space)
        self.resize(self.width(), height)
        self.setMinimumHeight(height)

    def clear(self):
        for item in self.items:
            item.setParent(None)
            item.destroy()
        del self.items[:]

    def resizeEvent(self, event):
        for i in range(len(self.items)):
            item = self.items[i]
            item.resize(self.width(), self.item_height)
            item.move(0, (self.item_height + self.item_space) * i)

class ProjectListTitle(QtGui.QWidget):
    def __init__(self, title, parent=None):
        super(ProjectListTitle, self).__init__(parent)
        self.title = title
        self.main_layout = QtGui.QHBoxLayout()
        self.main_layout.setContentsMargins(0, 0, 0, 0)
        self.setLayout(self.main_layout)

        self.label = QtGui.QLabel(self.title)
        self.main_layout.addWidget(self.label)
        self.main_layout.addStretch()
        self.sort_type = QtGui.QComboBox()
        self.sort_type.addItem(u"创建时间")
        self.sort_type.addItem(u"名字")
        self.main_layout.addWidget(self.sort_type)

class ProjectListArea(QtGui.QWidget):
    def __init__(self, parent=None):
        super(ProjectListArea, self).__init__(parent)
        self.tool_height = 40
        self.main_layout = QtGui.QVBoxLayout()
        self.setLayout(self.main_layout)

        self.button_layout = QtGui.QHBoxLayout()
        self.button_layout.setContentsMargins(0, 0, 0, 0)
        self.button_layout.setSpacing(0)

        self.search_box = QtGui.QLineEdit("Type to filter")
        self.search_box.setMinimumHeight(self.tool_height * 0.7)
        self.button_layout.addWidget(self.search_box)

        self.button_layout.addStretch()

        self.b_add_project = QtGui.QPushButton(u"创建项目")
        self.b_add_project.setMinimumHeight(self.tool_height)
        self.b_add_project.clicked.connect(self.create_project)
        self.button_layout.addWidget(self.b_add_project)

        self.main_layout.addLayout(self.button_layout)

        self.project_layout = QtGui.QGridLayout()
        self.main_layout.addLayout(self.project_layout)
        self.project_layout.setRowStretch(1, 1)

        self.recent_title = ProjectListTitle(u"最近的项目")
        self.recent_project_list = ProjectListView()
        self.recent_scroll_area = QtGui.QScrollArea()
        self.recent_scroll_area.setWidgetResizable(True)
        self.recent_scroll_area.setWidget(self.recent_project_list)

        self.my_title = ProjectListTitle(u"我的项目")
        self.my_project_list = ProjectListView()
        self.my_scroll_area = QtGui.QScrollArea()
        self.my_scroll_area.setWidgetResizable(True)
        self.my_scroll_area.setWidget(self.my_project_list)

        self.all_title = ProjectListTitle(u"所有项目")
        self.all_project_list = ProjectListView()
        self.all_scroll_area = QtGui.QScrollArea()
        self.all_scroll_area.setWidgetResizable(True)
        self.all_scroll_area.setWidget(self.all_project_list)

        self.project_layout.addWidget(self.recent_title, 0, 0)
        self.project_layout.addWidget(self.my_title, 0, 1)
        self.project_layout.addWidget(self.all_title, 0, 2)
        self.project_layout.addWidget(self.recent_scroll_area, 1, 0)
        self.project_layout.addWidget(self.my_scroll_area, 1, 1)
        self.project_layout.addWidget(self.all_scroll_area, 1, 2)

    def load_current(self):
        self.load_recent_project()
        self.load_my_project()
        self.load_active_project()

    def load_recent_project(self):
        self.recent_project_list.clear()
        recent_project = honey.Project.projects()
        for project in recent_project:
            self.recent_project_list.add_item(ProjectItem(project))

    def load_my_project(self):
        self.my_project_list.clear()
        my_project = honey.Project.projects()
        for project in my_project:
            self.my_project_list.add_item(ProjectItem(project))

    def load_active_project(self):
        self.all_project_list.clear()
        all_project = honey.Project.projects()
        for project in all_project:
            self.all_project_list.add_item(ProjectItem(project))

    def create_project(self):
        dialog = AddProjectDialog(self)

        if dialog.exec_() == QtGui.QDialog.Accepted:
            data = dialog.get_data()
            project = honey.Project(data)
            if project.add_to_db():
                hc.Project = project
                self.load_current()
                return True
        return False

class ProjectDynamics(QtGui.QWidget):
    def __init__(self, parent=None):
        super(ProjectDynamics, self).__init__(parent)
        self.main_layout = QtGui.QVBoxLayout()
        self.setLayout(self.main_layout)
        self.label = QtGui.QLabel(u"项目动态")
        self.label.setMinimumHeight(40)
        self.main_layout.addWidget(self.label)
        self.main_layout.addStretch()

class ProjectHome(QtGui.QSplitter):
    def __init__(self, parent=None):
        super(ProjectHome, self).__init__(parent)
        self.setContentsMargins(10, 10, 10, 10)

        self.left_area = ProjectListArea()
        self.right_area = ProjectDynamics()

        self.addWidget(self.left_area)
        self.addWidget(self.right_area)
        self.setSizes([self.width()*0.6, self.width()*0.4])

        self.load_current()

    def load_current(self):
        self.left_area.load_current()
        return True

if __name__ == '__main__':
    app = QtGui.QApplication(sys.argv)
    a = ProjectHome()
    a.resize(800, 700)
    a.show()
    a.load_current()
    sys.exit(app.exec_())


# Q1

def func1(N):
    answer=0
    for n in range(1,N+1):
        if n%3==0 or n%5==0:
            answer+=n
    return answer

N=int(input())
print(func1(N))


from sklearn.cluster import AgglomerativeClustering, MiniBatchKMeans, SpectralClustering, DBSCAN, AffinityPropagation
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from nltk.tokenize.treebank import TreebankWordDetokenizer
from sklearn.metrics.cluster import normalized_mutual_info_score
import pandas as pd
from gensim.models.ldamodel import LdaModel
import pickle
import numpy as np
from gensim.models.doc2vec import TaggedDocument, Doc2Vec
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
from sklearn.decomposition import TruncatedSVD
from collections import OrderedDict
from scipy.sparse import csr_matrix

def main():
    news_df = pd.read_pickle("news_df.pkl")
    no_of_topics = 20
    detokens = [TreebankWordDetokenizer().detokenize(token_list) for token_list in news_df['DocTokens']]
    bow_vectorizer = CountVectorizer()
    X_bow = bow_vectorizer.fit_transform(detokens)
    print(X_bow.shape)
    km_bow = AgglomerativeClustering(n_clusters=no_of_topics)
    y_bow = km_bow.fit_predict(X_bow.todense())
    print(len(y_bow))
    nmi_score = normalized_mutual_info_score(news_df['Topic'], y_bow)
    print(nmi_score)

    # ToDO t-SNE

    tfidf_vectorizer = TfidfVectorizer()
    X_tfidf = tfidf_vectorizer.fit_transform(detokens)
    print(X_tfidf.shape)
    km_tfidf = MiniBatchKMeans(n_clusters=no_of_topics)
    y_tfdif = km_tfidf.fit_predict(X_tfidf)
    print(len(y_tfdif))
    nmi_score = normalized_mutual_info_score(news_df['Topic'], y_tfdif)
    print(nmi_score)

    # tsne = TSNE(n_components=2)
    # #tsne = TruncatedSVD(n_components=2)
    # # tsne to our document vectors
    # tsne_tfidf = tsne.fit_transform(X_tfidf)
    # plt.figure(figsize=(10, 10))
    # plt.xticks(fontsize=12)
    # plt.yticks(fontsize=14)
    # plt.xlabel('tsne - 1', fontsize=20)
    # plt.ylabel('tsne - 2', fontsize=20)
    # plt.title("t-sne for Clustering TF IDF- 1", fontsize=20)
    # plt.scatter(tsne_tfidf[:, 0], tsne_tfidf[:, 1], c=y_tfdif, cmap=plt.cm.tab20)
    # plt.savefig('clustering_tfidf.png')
    # plt.show()

    lda_model_1 = LdaModel.load('ldamodel_voc1.model')
    bow_voc1_corpus = pickle.load(open('bow_voc1_corpus.pkl', 'rb'))
    predict_topics = []
    for i in range(len(bow_voc1_corpus)):
        topic = 0
        temp_pred = 0
        for (x, y) in lda_model_1[bow_voc1_corpus[i]]:
            if y > temp_pred:
                temp_pred = y
                topic = x
        predict_topics.append(topic)

    km_topic_dist = AgglomerativeClustering(n_clusters=no_of_topics)
    y_topic_dist = km_topic_dist.fit_predict(np.array(predict_topics).reshape(-1, 1))
    print(len(y_topic_dist))
    nmi_score = normalized_mutual_info_score(news_df['Topic'], y_topic_dist)
    print(nmi_score)

    top_dist = []
    keys = []
    for d in bow_voc1_corpus:
        tmp = {i: 0 for i in range(no_of_topics)}
        tmp.update(dict(lda_model_1[d]))
        vals = list(OrderedDict(tmp).values())
        top_dist += [vals]
        keys += [np.argmax(vals)]

    doc2vec_model_1 = Doc2Vec.load('doc2vec_voc1.model')
    X_doc2vec = doc2vec_model_1.docvecs.vectors_docs
    km_doc2vec = AgglomerativeClustering(n_clusters=no_of_topics)
    y_doc2vec = km_doc2vec.fit_predict(X_doc2vec)
    print(len(y_doc2vec))
    nmi_score = normalized_mutual_info_score(news_df['Topic'], y_doc2vec)
    print(nmi_score)

    # tsne = TSNE(n_components=2, perplexity=50, early_exaggeration=15, learning_rate=500, n_iter=2000)
    # X_tsne = tsne.fit_transform(X_doc2vec)
    # plt.figure(figsize=(10, 10))
    # plt.xticks(fontsize=12)
    # plt.yticks(fontsize=14)
    # plt.xlabel('tsne - 1', fontsize=20)
    # plt.ylabel('tsne - 2', fontsize=20)
    # plt.title("t-sne for LDA Model 1", fontsize=20)
    # plt.scatter(X_tsne[:, 0], X_tsne[:, 1], c=y_doc2vec, cmap=plt.cm.tab20)
    # plt.legend(loc='best')
    # plt.savefig('clustering_ldamodel1.png')
    # plt.show()

if __name__ == '__main__':
    main()